meta-llama/llama-2-7b-hf

app.py

@@ -1,41 +1,12 @@
-import gradio as gr
 from transformers import pipeline
-import torch  # <-- ADD THIS LINE
 
-# Load GPT-2 with optimized parameters
 model = pipeline(
     "text-generation",
-    model="gpt2",
-    max_length=200,
-    temperature=0.7,
-    early_stopping=True,
-    torch_dtype=torch.float32  # CPU compatibility
+    model="meta-llama/llama-2-7b-hf",
+    torch_dtype=torch.float16  # Reduce memory usage
 )
 
-def chat_with_gpt2(user_input):
-    prompt = f"{user_input}"
-    try:
-        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
-demo = gr.Interface(
-    fn=chat_with_gpt2,
-    inputs="text",
-    outputs="text",
-    title="GPT-2 Chat Interface",
-    description="Type any message to chat with GPT-2.",
-    flagging_mode="never"
-)
-
-# Launch the app
-if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        debug=True,
-        prevent_thread_lock=True
-    )
+def generate_test_cases(requirement):
+    prompt = f"Generate test cases for '{requirement}' in JSON format. Output only the array."
+    result = model(prompt, max_length=300)[0]["generated_text"]
+    return result.strip()