Upload model_utils.py

model/model_utils.py

@@ -1,19 +1,17 @@
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, T5ForConditionalGeneration
 import torch
 
 def load_model():
-    model_name = "mrm8488/codebert-base-finetuned-stackoverflow"
+    model_name = "Salesforce/codet5-small"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSequenceClassification.from_pretrained(model_name)
+    model = T5ForConditionalGeneration.from_pretrained(model_name)
     model.eval()
+    model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+    return tokenizer, model
 
-    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)"
+def generate_explanation(code, tokenizer, model):
+    device = model.device
+    input_text = "summarize: " + code
+    input_ids = tokenizer.encode(input_text, return_tensors="pt", truncation=True).to(device)
+    output = model.generate(input_ids, max_new_tokens=150, early_stopping=True)
+    return tokenizer.decode(output[0], skip_special_tokens=True)