distilgpt2 model pipeline..

app.py

@@ -1,35 +1,47 @@
 from transformers import pipeline
+import torch
 
-# Initialize the model pipeline
-print("Initializing model pipeline...")
+# Initialize the smaller GPT-2 model (distilgpt2)
+print("Initializing distilgpt2 model pipeline...")
 try:
     model = pipeline(
         "text-generation",
-        model="EleutherAI/gpt-neo-1.3B",  # Use GPT-Neo 1.3B as an example
-        torch_dtype=torch.float32  # Use float32 to avoid mixed precision issues
+        model="distilgpt2",  # Use the smaller, efficient distilgpt2 model
+        torch_dtype=torch.float32  # Use float32 for CPU compatibility
     )
-    print("Model pipeline initialized successfully.")
+    print("distilgpt2 model pipeline initialized successfully.")
 except Exception as e:
     print(f"Error initializing model: {e}")
     raise
 
 def generate_test_cases(requirement):
-    # Creating a prompt for generating test cases
+    # Formulate the prompt to ask the model to generate test cases in JSON format
     prompt = f"Generate test cases for the following software requirement in JSON format: '{requirement}'. Only provide the JSON array of test cases."
     
     try:
         print("Generating test cases...")
+        # Generate text (test cases) with a max length of 300 characters
         result = model(prompt, max_length=300, num_return_sequences=1)[0]["generated_text"]
         print("Test cases generated successfully.")
         return result.strip()
     except Exception as e:
         print(f"Error during generation: {e}")
-        raise
+        return None
 
 # Example usage
 if __name__ == "__main__":
-    requirement = "User login functionality"
-    try:
-        print(generate_test_cases(requirement))
-    except Exception as e:
-        print(f"Error in application: {e}")
+    # Sample requirements for testing
+    requirements = [
+        "User login functionality with email and password",
+        "Search functionality on an e-commerce website",
+        "Order placement process for an online store",
+        "Admin dashboard to manage users and content"
+    ]
+    
+    for req in requirements:
+        print(f"\nTest cases for requirement: {req}")
+        test_cases = generate_test_cases(req)
+        if test_cases:
+            print(test_cases)
+        else:
+            print("Failed to generate test cases.")