Spaces:
Sleeping
Sleeping
File size: 828 Bytes
841ffa5 2cfa511 841ffa5 2cfa511 841ffa5 2cfa511 841ffa5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
def load_model():
model_name = "mrm8488/codebert-base-finetuned-stackoverflow"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
model.eval()
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
return tokenizer, model, device
def generate_explanation(code, tokenizer, model, device):
inputs = tokenizer(code, return_tensors="pt", truncation=True, padding=True).to(device)
with torch.no_grad():
logits = model(**inputs).logits
predicted_class_id = logits.argmax().item()
return f"This code is classified as category ID: {predicted_class_id} (label may vary based on fine-tuning objective)"
|