Update app.py

app.py

@@ -1,50 +1,33 @@
 import gradio as gr
 from transformers import pipeline
-import torch
 
 # Load GPT-2 with optimized parameters
 model = pipeline(
     "text-generation",
     model="gpt2",
-    max_length=300,          # Increased length to allow more detailed output
-    temperature=0.7,         # Balance between creativity and stability
-    early_stopping=True,     # Stop generation early if completed
+    max_length=200,          # Maximum response length
+    temperature=0.7,         # Creativity (0.0 = deterministic, 1.0 = random)
+    early_stopping=True,     # Stop generation early if possible
     torch_dtype=torch.float32  # CPU compatibility
 )
 
-def generate_test_cases(requirement):
-    # Revised prompt to focus on the input requirement and avoid copying the example
-    prompt = f"""
-    Generate test cases for the requirement: '{requirement}' in JSON array format. 
-    Each test case must include "id", "title", "steps" (array of strings), and "expected_result". 
-
-    DO NOT use the example below in your response. Output only the JSON array, no explanations.
-
-    Example format (IGNORE this when generating):
-    [
-        {{
-            "id": 1,
-            "title": "Example Test",
-            "steps": ["Step 1", "Step 2"],
-            "expected_result": "Expected outcome"
-        }}
-    ]
-    """
-    
+def chat_with_gpt2(user_input):
+    # Generate a response based on user input
+    prompt = f"{user_input}"
     try:
         with torch.no_grad():
-            result = model(prompt, max_time=15)[0]["generated_text"]
+            result = model(prompt, max_time=10)[0]["generated_text"]
         return result.strip()
     except Exception as e:
         return f"Error: {str(e)}"
 
 # Create Gradio interface
 demo = gr.Interface(
-    fn=generate_test_cases,
+    fn=chat_with_gpt2,
     inputs="text",
     outputs="text",
-    title="Test Case Generator",
-    description="Enter a requirement to generate test cases.",
+    title="GPT-2 Chat Interface",
+    description="Type any message to chat with GPT-2.",
     flagging_mode="never"
 )