gpt neo

app.py

@@ -1,13 +1,20 @@
 from transformers import pipeline
 import torch  # Import torch library
 
+# Use the GPT-Neo model, a larger and more capable alternative to GPT-2
 model = pipeline(
     "text-generation",
-    model="meta-llama/llama-2-7b-hf",
+    model="EleutherAI/gpt-neo-1.3B",  # GPT-Neo 1.3B model
     torch_dtype=torch.float16  # Reduce memory usage
 )
 
 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"]
+    # Creating a prompt for generating test cases
+    prompt = f"Generate test cases for the following software requirement in JSON format: '{requirement}'. Only provide the JSON array of test cases."
+    result = model(prompt, max_length=300, num_return_sequences=1)[0]["generated_text"]
     return result.strip()
+
+# Example usage
+if __name__ == "__main__":
+    requirement = "User login functionality"
+    print(generate_test_cases(requirement))