Update app.py

app.py

@@ -3,12 +3,12 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 @st.cache_resource
 def load_model():
-    model_name = "Salesforce/codet5-small"
+    model_name = "Salesforce/codet5-base"  # Switch to 'codet5-base' for better results
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
     return tokenizer, model
 
-# Load model
+# Load the model and tokenizer
 tokenizer, model = load_model()
 
 st.title("Code Generator")
@@ -21,14 +21,25 @@ if st.button("Generate Code"):
     if prompt.strip():
         with st.spinner("Generating code..."):
             inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True)
-            outputs = model.generate(inputs.input_ids, max_length=max_length, num_beams=5, temperature=0.7, early_stopping=True)
 
+            # Use sampling-based generation for better quality
+            outputs = model.generate(
+                inputs.input_ids,
+                max_length=max_length,
+                temperature=0.7,
+                top_p=0.95,
+                do_sample=True,
+            )
+
+            # Debugging: Show raw token output
             st.write("### Debugging: Raw Model Output")
-            st.json(outputs.tolist())  # Debugging output
+            st.json(outputs.tolist())
+
+            # Decode tokens properly
+            generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
 
-            generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
-            
             st.write("### Generated Code:")
             st.code(generated_code, language="python")
+
     else:
         st.warning("Please enter a prompt!")