newwwww

model/generate.py

@@ -82,55 +82,33 @@ def get_optimal_model_for_memory():
     available_memory = psutil.virtual_memory().available / (1024 * 1024)  # MB
     logger.info(f"Available memory: {available_memory:.1f}MB")
 
-    if available_memory < 500:
+    if available_memory < 300:
         return None  # Use template fallback
-    elif available_memory < 1000:
+    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 - FIXED VERSION."""
+    """Load model with low memory settings."""
     try:
         logger.info(f"Loading {model_name} with memory optimizations...")
 
-        # Load tokenizer first
-        tokenizer = AutoTokenizer.from_pretrained(
-            model_name, 
-            padding_side='left', 
-            use_fast=True
-        )
+        tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side='left', use_fast=True)
 
-        # Fix tokenizer pad token issue
         if tokenizer.pad_token is None:
-            if tokenizer.eos_token is not None:
-                tokenizer.pad_token = tokenizer.eos_token
-            else:
-                tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+            tokenizer.pad_token = tokenizer.eos_token
 
-        # Load model with corrected parameters
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            device_map=None,  # Don't use device_map with low_cpu_mem_usage
-            low_cpu_mem_usage=False,  # Disable to avoid meta tensor issues
-            use_cache=True,  # Enable cache for better performance
-            trust_remote_code=True
+            torch_dtype=torch.float16,
+            device_map="cpu",
+            low_cpu_mem_usage=True,
+            use_cache=False,
         )
 
-        # Move to CPU explicitly if needed
-        if torch.cuda.is_available():
-            model = model.to('cpu')
-        
         model.eval()
-        
-        # Only enable gradient checkpointing if model supports it
-        if hasattr(model, 'gradient_checkpointing_enable'):
-            try:
-                model.gradient_checkpointing_enable()
-            except Exception as e:
-                logger.warning(f"Could not enable gradient checkpointing: {e}")
-        
+        model.gradient_checkpointing_enable()
         logger.info(f"✅ Model {model_name} loaded successfully")
         return tokenizer, model
 
@@ -190,20 +168,6 @@ def generate_authentication_tests(matches: List[str]) -> List[Dict]:
             ],
             "expected": "Login fails with appropriate error message 'Invalid credentials'",
             "postconditions": ["User remains on login page", "Account security maintained"],
-            "test_data": "Invalid username: invalid@example.com, Valid password: Test@123"
-        },
-        {
-            "title": "Invalid Password Login",
-            "description": "Verify that login fails with invalid password",
-            "preconditions": ["Application is accessible"],
-            "steps": [
-                "Navigate to login page",
-                "Enter valid username",
-                "Enter invalid password",
-                "Click login button"
-            ],
-            "expected": "Login fails with appropriate error message 'Invalid credentials'",
-            "postconditions": ["User remains on login page", "Account security maintained"],
             "test_data": "Valid username: testuser@example.com, Invalid password: WrongPass123"
         },
         {
@@ -595,7 +559,7 @@ def parse_generated_test_cases(generated_text: str) -> List[Dict]:
     return test_cases
 
 def generate_with_ai_model(srs_text: str, tokenizer, model) -> List[Dict]:
-    """Generate test cases using AI model - FIXED VERSION"""
+    """Generate test cases using AI model"""
     max_input_length = 300
     if len(srs_text) > max_input_length:
         srs_text = srs_text[:max_input_length]
@@ -607,40 +571,27 @@ Test Cases:
 1."""
 
     try:
-        # Encode with proper attention mask
-        inputs = tokenizer(
+        inputs = tokenizer.encode(
             prompt,
             return_tensors="pt",
             max_length=200,
-            truncation=True,
-            padding=True,
-            return_attention_mask=True
+            truncation=True
         )
 
-        # Extract input_ids and attention_mask
-        input_ids = inputs['input_ids']
-        attention_mask = inputs['attention_mask']
-
         with torch.no_grad():
             outputs = model.generate(
-                input_ids,
-                attention_mask=attention_mask,  # Provide attention mask
+                inputs,
                 max_new_tokens=150,
                 num_return_sequences=1,
                 temperature=0.7,
                 do_sample=True,
-                pad_token_id=tokenizer.pad_token_id,
-                eos_token_id=tokenizer.eos_token_id,
-                use_cache=True,
+                pad_token_id=tokenizer.eos_token_id,
+                use_cache=False,
             )
 
         generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Clean up tensors
-        del inputs, input_ids, attention_mask, outputs
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        
+        del inputs, outputs
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None
         return parse_generated_test_cases(generated_text)
 
     except Exception as e:
@@ -729,4 +680,5 @@ def get_algorithm_reason(model_name: str) -> str:
                 "Uses comprehensive pattern matching, requirement analysis, and structured test case templates for robust test coverage.")
     else:
         return ("Model selected based on optimal tradeoff between memory usage, language generation capability, "
-                "and test case quality requirements.")
+                "and test case quality requirements.")
+