Views
No views yet
1import torch
2from sentence_transformers import SentenceTransformer
3from transformers import AutoTokenizer
4
5# Load tokenizer
6tokenizer = AutoTokenizer.from_pretrained("nomic-ai/CodeRankEmbed", trust_remote_code=True)
7
8# Load your trained model
9# (You'll need to load the SimilarityModel class and state dict)
10
11# Example inference
12code1 = "..."
13code2 = "..."
14
15# Tokenize
16enc1 = tokenizer(code1, return_tensors='pt', padding=True, truncation=True, max_length=512)
17enc2 = tokenizer(code2, return_tensors='pt', padding=True, truncation=True, max_length=512)
18
19# Compute similarity
20with torch.no_grad():
21 similarity = model(enc1['input_ids'], enc1['attention_mask'],
22 enc2['input_ids'], enc2['attention_mask'])
23
24print(f"Similarity: {similarity.item():.3f}")
25# > 0.7: High similarity (likely to pass)
26# < 0.3: Low similarity (likely to fail)1@misc{coderankembed2024,
2 title={CodeRankEmbed: Code Similarity Learning},
3 author={Nomic AI},
4 year={2024},
5 url={https://huggingface.co/nomic-ai/CodeRankEmbed}
6}