newwww1w3

model/generate.py

@@ -15,31 +15,26 @@ MEMORY_OPTIMIZED_MODELS = [
     "gpt2",  # ~500MB
     "distilgpt2",  # ~250MB
     "microsoft/DialoGPT-small",  # ~250MB
-    "huggingface/CodeBERTa-small-v1",  # Code tasks
 ]
 
-# Singleton state
 _generator_instance = None
 
 def get_optimal_model_for_memory():
-    """Select the best model based on available memory."""
     available_memory = psutil.virtual_memory().available / (1024 * 1024)  # MB
     logger.info(f"Available memory: {available_memory:.1f}MB")
 
     if available_memory < 300:
-        return None  # Use template fallback
+        return None
     elif available_memory < 600:
         return "microsoft/DialoGPT-small"
     else:
         return "distilgpt2"
 
 def load_model_with_memory_optimization(model_name):
-    """Load model with low memory settings."""
     try:
         logger.info(f"Loading {model_name} with memory optimizations...")
 
         tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side='left', use_fast=True)
-
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
 
@@ -72,104 +67,119 @@ def extract_keywords(text):
 def generate_template_based_test_cases(srs_text):
     keywords = extract_keywords(srs_text)
     test_cases = []
+    counter = 1
 
     if any(word in keywords for word in ['login', 'authentication', 'user', 'password']):
         test_cases.extend([
             {
-                "id": "TC_001",
+                "id": f"TC_{counter:03d}",
                 "title": "Valid Login Test",
                 "description": "Test login with valid credentials",
                 "steps": ["Enter valid username", "Enter valid password", "Click login"],
                 "expected": "User should be logged in successfully"
             },
             {
-                "id": "TC_002",
+                "id": f"TC_{counter+1:03d}",
                 "title": "Invalid Login Test",
                 "description": "Test login with invalid credentials",
                 "steps": ["Enter invalid username", "Enter invalid password", "Click login"],
                 "expected": "Error message should be displayed"
             }
         ])
+        counter += 2
 
     if any(word in keywords for word in ['database', 'data', 'store', 'save']):
         test_cases.append({
-            "id": "TC_003",
+            "id": f"TC_{counter:03d}",
             "title": "Data Storage Test",
             "description": "Test data storage functionality",
             "steps": ["Enter data", "Save data", "Verify storage"],
             "expected": "Data should be stored correctly"
         })
+        counter += 1
+
+    if any(word in keywords for word in ['validation', 'error']):
+        test_cases.append({
+            "id": f"TC_{counter:03d}",
+            "title": "Input Validation Test",
+            "description": "Test system input validation",
+            "steps": ["Enter invalid input", "Submit form"],
+            "expected": "System should prevent submission and show error"
+        })
 
     if not test_cases:
-        test_cases = [
-            {
-                "id": "TC_001",
-                "title": "Basic Functionality Test",
-                "description": "Test basic system functionality",
-                "steps": ["Access the system", "Perform basic operations", "Verify results"],
-                "expected": "System should work as expected"
-            }
-        ]
+        test_cases = [{
+            "id": "TC_001",
+            "title": "Generic Functional Test",
+            "description": "Test basic system functionality",
+            "steps": ["Access system", "Perform operations"],
+            "expected": "System works correctly"
+        }]
 
     return test_cases
 
-def parse_generated_test_cases(generated_text):
-    lines = generated_text.split('\n')
+def parse_generated_test_cases(text):
+    lines = text.split('\n')
     test_cases = []
-    current_case = {}
+    current = {}
+    steps = []
     case_counter = 1
 
     for line in lines:
         line = line.strip()
-        if line.startswith(('1.', '2.', '3.', 'TC', 'Test')):
-            if current_case:
-                test_cases.append(current_case)
-            current_case = {
+        if re.match(r'^\d+\.', line) or line.lower().startswith("test case"):
+            if current:
+                current["steps"] = steps or ["Execute the test"]
+                current["expected"] = "Test should pass"
+                test_cases.append(current)
+            current = {
                 "id": f"TC_{case_counter:03d}",
                 "title": line,
-                "description": line,
-                "steps": ["Execute the test"],
-                "expected": "Test should pass"
+                "description": line
             }
+            steps = []
             case_counter += 1
+        elif line.lower().startswith("step") or line.startswith("-"):
+            steps.append(line.lstrip('- ').strip())
 
-    if current_case:
-        test_cases.append(current_case)
+    if current:
+        current["steps"] = steps or ["Execute the test"]
+        current["expected"] = "Test should pass"
+        test_cases.append(current)
 
     if not test_cases:
         return [{
             "id": "TC_001",
             "title": "Generated Test Case",
-            "description": "Auto-generated test case based on requirements",
-            "steps": ["Review requirements", "Execute test", "Verify results"],
-            "expected": "Requirements should be met"
+            "description": "Auto-generated based on SRS",
+            "steps": ["Review requirements", "Execute test"],
+            "expected": "Requirements met"
         }]
 
     return test_cases
 
 def generate_with_ai_model(srs_text, tokenizer, model):
-    max_input_length = 200
-    if len(srs_text) > max_input_length:
-        srs_text = srs_text[:max_input_length]
-
-    prompt = f"""Generate test cases for this software requirement:
+    prompt = f"""Generate detailed and numbered test cases for the following software requirement:
 {srs_text}
 
 Test Cases:
 1."""
 
+    input_length = len(srs_text.split())
+    max_new_tokens = min(max(100, input_length * 2), 600)
+
     try:
         inputs = tokenizer.encode(
             prompt,
             return_tensors="pt",
-            max_length=150,
-            truncation=True
+            truncation=True,
+            max_length=512
         )
 
         with torch.no_grad():
             outputs = model.generate(
                 inputs,
-                max_new_tokens=100,
+                max_new_tokens=max_new_tokens,
                 num_return_sequences=1,
                 temperature=0.7,
                 do_sample=True,
@@ -203,32 +213,38 @@ def generate_with_fallback(srs_text):
     test_cases = generate_template_based_test_cases(srs_text)
     return test_cases, "Template-Based Generator", "rule-based", "Low memory - fallback to rule-based generation"
 
-# ✅ Function exposed to app.py
 def generate_test_cases(srs_text):
     return generate_with_fallback(srs_text)[0]
 
+def generate_test_cases_and_info(input_text):
+    test_cases, model_name, algorithm_used, reason = generate_with_fallback(input_text)
+    return {
+        "model": model_name,
+        "algorithm": algorithm_used,
+        "reason": reason,
+        "test_cases": test_cases
+    }
+
 def get_generator():
     global _generator_instance
     if _generator_instance is None:
         class Generator:
             def __init__(self):
                 self.model_name = get_optimal_model_for_memory()
-                self.tokenizer = None
-                self.model = None
+                self.tokenizer, self.model = None, None
                 if self.model_name:
                     self.tokenizer, self.model = load_model_with_memory_optimization(self.model_name)
 
             def get_model_info(self):
                 mem = psutil.Process().memory_info().rss / 1024 / 1024
                 return {
-                    "model_name": self.model_name if self.model_name else "Template-Based Generator",
+                    "model_name": self.model_name or "Template-Based Generator",
                     "status": "loaded" if self.model else "template_mode",
                     "memory_usage": f"{mem:.1f}MB",
                     "optimization": "low_memory"
                 }
 
         _generator_instance = Generator()
-
     return _generator_instance
 
 def monitor_memory():
@@ -238,25 +254,14 @@ def monitor_memory():
         gc.collect()
         logger.info("Memory cleanup triggered")
 
-# ✅ NEW FUNCTION for enhanced output: test cases + model info + reason
-def generate_test_cases_and_info(input_text):
-    test_cases, model_name, algorithm_used, reason = generate_with_fallback(input_text)
-    return {
-        "model": model_name,
-        "algorithm": algorithm_used,
-        "reason": reason,
-        "test_cases": test_cases
-    }
-
-# ✅ Explain why each algorithm is selected
 def get_algorithm_reason(model_name):
     if model_name == "microsoft/DialoGPT-small":
         return "Selected due to low memory availability; DialoGPT-small provides conversational understanding in limited memory environments."
     elif model_name == "distilgpt2":
-        return "Selected for its balance between performance and low memory usage. Ideal for small environments needing causal language modeling."
+        return "Selected for its balance between performance and low memory usage."
     elif model_name == "gpt2":
-        return "Chosen for general-purpose text generation with moderate memory headroom."
+        return "Chosen for general-purpose generation with moderate memory headroom."
     elif model_name is None:
-        return "No model used due to insufficient memory. Rule-based template generation chosen instead."
+        return "Rule-based fallback due to memory constraints."
     else:
-        return "Model selected based on best tradeoff between memory usage and language generation capability."
+        return "Chosen based on available memory and task compatibility."