Update app.py

app.py

@@ -1,18 +1,18 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 
-# Model aur tokenizer load karte hain
-model_name = "huggingface/CodeParrot-small"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+# Load CodeT5 model and tokenizer
+model_name = "Salesforce/codet5-base"
+tokenizer = T5Tokenizer.from_pretrained(model_name)
+model = T5ForConditionalGeneration.from_pretrained(model_name)
 
 def generate_comments(code):
-    inputs = tokenizer.encode(code, return_tensors="pt")
-    outputs = model.generate(inputs, max_length=200, num_return_sequences=1)
+    inputs = tokenizer("Comment this code: " + code, return_tensors="pt")
+    outputs = model.generate(inputs["input_ids"], max_length=200, num_return_sequences=1)
     commented_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return commented_code
 
-# Gradio interface banate hain
+# Create Gradio interface
 iface = gr.Interface(
     fn=generate_comments,
     inputs="text",