Upload model_utils.py

model/model_utils.py

@@ -1,22 +1,19 @@
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 
 def load_model():
-    model_name = "Salesforce/codet5-base"
+    model_name = "mrm8488/codebert-base-finetuned-stackoverflow"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSeq2SeqLM.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(prompt, tokenizer, model, device):
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True).to(device)
-    output = model.generate(
-        **inputs,
-        decoder_start_token_id=tokenizer.pad_token_id,
-        max_new_tokens=256,
-        temperature=0.7
-    )
-    return tokenizer.decode(output[0], skip_special_tokens=True)
+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)"