Update app.py

app.py

@@ -3,10 +3,10 @@ from transformers import pipeline
 import json
 import signal
 
-# Load a small, free, instruction-following model
-generator = pipeline("text2text-generation", model="google/flan-t5-base")
+# Load a small, instruction-following model
+generator = pipeline("text2text-generation", model="google/flan-t5-large")
 
-# Function to handle timeouts
+# Timeout handling
 class TimeoutException(Exception):
     pass
 
@@ -17,36 +17,33 @@ signal.signal(signal.SIGALRM, timeout_handler)
 
 def generate_test_cases(user_story):
     try:
-        # Set a 3-minute timeout
-        signal.alarm(180)
+        signal.alarm(180)  # Set a 3-minute timeout
 
-        # Force JSON output using a structured prompt
+        # Improved structured prompt
         prompt = (
-            f"Convert the following user story into structured JSON test cases.\n"
-            f"User story: {user_story}\n"
-            f"Output format:\n"
-            f'{{"test_cases": [{{"id": 1, "title": "Test case title", "steps": ["Step 1", "Step 2"], "expected_result": "Expected outcome"}}]}}'
+            f"Given the user story below, generate software test cases in JSON format.\n"
+            f"User Story: {user_story}\n"
+            f"Format:\n"
+            f'{{"test_cases": [{{"id": 1, "title": "Test case title", "steps": ["Step 1", "Step 2"], "expected_result": "Expected outcome"}}]}}\n'
+            f"Ensure the output follows the JSON format strictly."
         )
 
-        # Generate output
         output = generator(prompt, max_length=512, do_sample=False)[0]["generated_text"]
 
-        # Try parsing as JSON to ensure proper formatting
+        # Attempt to parse JSON
         try:
             json_output = json.loads(output)
             formatted_output = json.dumps(json_output, indent=4)
         except json.JSONDecodeError:
             formatted_output = "Error: Model did not return valid JSON. Try re-running."
 
-        # Disable the alarm if processing is successful
-        signal.alarm(0)
-
+        signal.alarm(0)  # Disable timeout if successful
         return formatted_output
 
     except TimeoutException:
         return "Processing timed out. Please try again with a simpler input."
 
-# Create Gradio interface
+# Gradio UI
 iface = gr.Interface(
     fn=generate_test_cases,
     inputs=gr.Textbox(lines=5, placeholder="Enter your user story here..."),