Spaces:
Sleeping
Sleeping
File size: 2,129 Bytes
677f286 595955c 677f286 595955c 677f286 595955c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | import os
import torch
import pickle
import numpy as np
from transformers import AutoTokenizer, AutoModel
# FORCE SMALL MODEL
MODEL_NAME = "huggingface/CodeBERTa-small-v1"
EMBEDDING_DIR = "reference_embeddings"
REFERENCE_CODE = {
"Python": "def factorial(n): return 1 if n == 0 else n * factorial(n-1)",
"C": "int factorial(int n) { return (n == 0) ? 1 : n * factorial(n - 1); }",
"C++": "int factorial(int n) { return (n == 0) ? 1 : n * factorial(n - 1); }",
"Java": "public int factorial(int n) { return (n == 0) ? 1 : n * factorial(n - 1); }",
"JavaScript": "function factorial(n) { return (n === 0) ? 1 : n * factorial(n - 1); }",
"Go": "func factorial(n int) int { if n == 0 { return 1 } return n * factorial(n-1) }",
"Rust": "fn factorial(n: u32) -> u32 { if n == 0 { 1 } else { n * factorial(n - 1) } }",
"PHP": "function factorial($n) { return ($n == 0) ? 1 : $n * factorial($n - 1); }",
"Ruby": "def factorial(n) return 1 if n == 0; n * factorial(n - 1) end",
"Swift": "func factorial(_ n: Int) -> Int { return (n == 0) ? 1 : n * factorial(n - 1) }",
"Kotlin": "fun factorial(n: Int): Int { return if (n == 0) 1 else n * factorial(n - 1) }",
}
def generate_embeddings():
print(f"Loading model: {MODEL_NAME}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModel.from_pretrained(MODEL_NAME)
if not os.path.exists(EMBEDDING_DIR):
os.makedirs(EMBEDDING_DIR)
print("Generating reference embeddings...")
for language, code in REFERENCE_CODE.items():
inputs = tokenizer(code, return_tensors="pt", truncation=True, max_length=512)
with torch.no_grad():
outputs = model(**inputs)
# Mean pooling to get 768 dimensions
embedding = outputs.last_hidden_state.mean(dim=1).squeeze().numpy()
file_path = os.path.join(EMBEDDING_DIR, f"{language}.pkl")
with open(file_path, "wb") as f:
pickle.dump(embedding, f)
print(f"Saved {language} embedding (Shape: {embedding.shape}).")
if __name__ == "__main__":
generate_embeddings()
|