Update app.py

app.py

@@ -1,36 +1,65 @@
-from transformers import pipeline
 import torch
+from transformers import pipeline, set_seed
+import time
 
-# Initialize the smaller GPT-2 model (distilgpt2)
+# Set random seed for reproducibility
+set_seed(42)
+
+# Initialize model pipeline with better error handling
 print("Initializing distilgpt2 model pipeline...")
-try:
-    model = pipeline(
-        "text-generation",
-        model="distilgpt2",  # Use the smaller, efficient distilgpt2 model
-        torch_dtype=torch.float32  # Use float32 for CPU compatibility
-    )
-    print("distilgpt2 model pipeline initialized successfully.")
-except Exception as e:
-    print(f"Error initializing model: {e}")
-    raise
 
+def initialize_model():
+    try:
+        # Set up the model pipeline with distilgpt2
+        model = pipeline(
+            "text-generation",
+            model="distilgpt2",  # Load the distilgpt2 model
+            tokenizer="distilgpt2",  # Ensure tokenizer is loaded as well
+            torch_dtype=torch.float32  # Ensures compatibility with CPU
+        )
+        print("Model pipeline initialized successfully.")
+        return model
+    except Exception as e:
+        print(f"Error during model initialization: {e}")
+        return None
+
+# Retry logic to handle initialization issues
+model = None
+attempts = 0
+max_attempts = 3
+while model is None and attempts < max_attempts:
+    print(f"Attempt {attempts + 1}/{max_attempts} to initialize the model...")
+    model = initialize_model()
+    attempts += 1
+    if model is None:
+        print(f"Retrying after a delay...")
+        time.sleep(10)  # Wait 10 seconds before retrying
+
+if model is None:
+    print("Failed to initialize the model after several attempts. Exiting.")
+    exit(1)  # Exit if model is not initialized successfully
+
+# Generate test cases function
 def generate_test_cases(requirement):
-    # 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."
+    prompt = f"Generate detailed test cases in JSON format for the following software requirement: '{requirement}'. Please return only the JSON array with 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"]
+        result = model(prompt, max_length=500, num_return_sequences=1)[0]["generated_text"]
         print("Test cases generated successfully.")
-        return result.strip()
+        
+        # Clean up output to only return JSON array
+        clean_result = result.strip().split("\n")[0]
+        if clean_result.startswith('[') and clean_result.endswith(']'):
+            return clean_result
+        else:
+            return f"Error: Unexpected output format.\n{result.strip()}"
     except Exception as e:
         print(f"Error during generation: {e}")
         return None
 
 # Example usage
 if __name__ == "__main__":
-    # Sample requirements for testing
     requirements = [
         "User login functionality with email and password",
         "Search functionality on an e-commerce website",