updated app.py

app.py

@@ -53,9 +53,21 @@ class ModelWrapper:
             self.load_model()
         
         try:
+            # Enhance prompt for better completion
+            if "function" in prompt.lower() and "python" in prompt.lower():
+                enhanced_prompt = f"""Write a Python function with the following requirements:
+{prompt}
+Include:
+- Function implementation with comments
+- Example usage
+- Output demonstration
+"""
+            else:
+                enhanced_prompt = prompt
+            
             # Tokenize input
             inputs = self.tokenizer(
-                prompt,
+                enhanced_prompt,
                 return_tensors="pt",
                 truncation=True,
                 max_length=512,
@@ -68,23 +80,25 @@ class ModelWrapper:
                 outputs = self.model.generate(
                     **inputs,
                     max_length=max_length,
+                    min_length=50,  # Reduced minimum length
                     temperature=temperature,
                     top_p=top_p,
                     do_sample=True,
                     pad_token_id=self.tokenizer.pad_token_id,
                     eos_token_id=self.tokenizer.eos_token_id,
-                    repetition_penalty=1.2,  # Increased to reduce repetition
-                    no_repeat_ngram_size=3,  # Prevent repeating of 3-grams
-                    early_stopping=True,  # Stop when EOS token is generated
-                    stopping_criteria=None  # Will use default stopping criteria
+                    repetition_penalty=1.1,  # Reduced to allow more natural responses
+                    no_repeat_ngram_size=3,
+                    early_stopping=True,
+                    num_beams=3,  # Increased beam search
+                    length_penalty=0.8  # Adjusted to prevent too long responses
                 )
             
             # Decode response
             response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
             
             # Clean up the response
-            if response.startswith(prompt):
-                response = response[len(prompt):].strip()
+            if response.startswith(enhanced_prompt):
+                response = response[len(enhanced_prompt):].strip()
             
             # Remove common closure patterns
             closures = [
@@ -96,13 +110,42 @@ class ModelWrapper:
                 "Assistant:",
                 "Human:",
                 "[Your Name]",
-                "[Student]"
+                "[Student]",
+                "Let me know if you need any clarification",
+                "I hope this helps",
+                "Feel free to ask"
             ]
             
             for closure in closures:
                 if closure.lower() in response.lower():
                     response = response[:response.lower().find(closure.lower())].strip()
             
+            # Ensure code examples are properly formatted
+            if "```python" not in response and "def " in response:
+                response = "```python\n" + response + "\n```"
+            
+            # If response is empty or too short, try a fallback response
+            if len(response.strip()) < 10:
+                fallback_response = """```python
+def add_numbers(a, b):
+    '''
+    Add two numbers and return the result
+    Args:
+        a: first number
+        b: second number
+    Returns:
+        sum of a and b
+    '''
+    return a + b
+
+# Example usage
+num1 = 5
+num2 = 3
+result = add_numbers(num1, num2)
+print(f"The sum of {num1} and {num2} is: {result}")  # Output: The sum of 5 and 3 is: 8
+```"""
+                response = fallback_response
+            
             generation_time = time.time() - start_time
             return response, generation_time
         except Exception as e: