Update app.py

app.py

@@ -1,20 +1,62 @@
 import gradio as gr
 from transformers import pipeline
+import torch
 
-model = pipeline("text-generation", model="gpt2")
+# Load GPT-2 with optimized parameters
+model = pipeline(
+    "text-generation",
+    model="gpt2",
+    max_length=200,
+    temperature=0.7,
+    early_stopping=True,
+    torch_dtype=torch.float16  # Use half-precision for speed
+)
 
 def generate_test_cases(requirement):
-    prompt = f"Generate test cases for '{requirement}' in JSON format. Output only the array."
-    result = model(prompt, max_length=200)[0]["generated_text"]
-    return result.strip()
+    # Simplified prompt with clear instructions and example
+    prompt = f"""
+    Generate test cases for '{requirement}' in JSON format. Output only the array.
+    
+    Example format:
+    [
+        {{
+            "id": 1,
+            "title": "Valid Login",
+            "steps": ["Enter valid email", "Enter valid password", "Click login"],
+            "expected_result": "Redirect to dashboard"
+        }},
+        {{
+            "id": 2,
+            "title": "Invalid Password",
+            "steps": ["Enter valid email", "Enter wrong password", "Click login"],
+            "expected_result": "Error message displayed"
+        }}
+    ]
+    """
+    
+    try:
+        # Generate text with a timeout (avoids long inference)
+        with torch.no_grad():
+            result = model(prompt, max_time=10)[0]["generated_text"]
+        return result.strip()
+    except Exception as e:
+        return f"Error: {str(e)}"
 
+# Create Gradio interface with proper flagging_mode
 demo = gr.Interface(
     fn=generate_test_cases,
     inputs="text",
     outputs="text",
     title="Test Case Generator",
     description="Enter a requirement to generate test cases.",
-    flagging_mode="never"  # Valid value for flagging_mode
+    flagging_mode="never"  # Valid value
 )
 
-demo.launch(server_name="0.0.0.0", server_port=7860)
+# Launch with proper settings
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        debug=True,
+        prevent_thread_lock=True
+    )