Update app.py

app.py

@@ -1,17 +1,36 @@
 import gradio as gr
 from transformers import pipeline
+import signal
 
-# Load an open-source model for free-tier compatibility
+# Load model
 generator = pipeline("text-generation", model="google/flan-t5-base")
 
+# Function to handle timeouts
+class TimeoutException(Exception):
+    pass
+
+def timeout_handler(signum, frame):
+    raise TimeoutException("Processing took too long. Try a simpler input.")
+
+signal.signal(signal.SIGALRM, timeout_handler)
+
 def generate_test_cases(user_story):
-    prompt = f"Generate structured test cases in JSON format from this user story: {user_story}"
-    
-    output = generator(prompt, max_length=300, do_sample=True)[0]["generated_text"]
+    try:
+        # Set a 3-minute timeout
+        signal.alarm(120)
+
+        prompt = f"Generate structured test cases in JSON format from this user story: {user_story}"
+        output = generator(prompt, max_length=300, do_sample=True)[0]["generated_text"]
 
-    return output
+        # Disable the alarm if processing is successful
+        signal.alarm(0)
 
-# Create a Gradio interface
+        return output
+
+    except TimeoutException:
+        return "Processing timed out. Please try again with a simpler input."
+
+# Create Gradio interface
 iface = gr.Interface(
     fn=generate_test_cases,
     inputs=gr.Textbox(lines=5, placeholder="Enter your user story here..."),
@@ -20,6 +39,5 @@ iface = gr.Interface(
     description="Enter a user story and get structured test cases in JSON format.",
 )
 
-
 if __name__ == "__main__":
     iface.launch()