newwww1w

model/generate.py

@@ -5,411 +5,258 @@ import logging
 import psutil
 import re
 import gc
-from typing import List, Dict, Any, Optional, Tuple
-from dataclasses import dataclass
 
-# Configure logging
-logging.basicConfig(level=logging.INFO)
+# Initialize logger
 logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
 
-# Constants
+# List of memory-optimized models
 MEMORY_OPTIMIZED_MODELS = [
-    "distilgpt2",  # ~250MB - default first choice
-    "microsoft/DialoGPT-small",  # ~250MB - good for conversational context
-    "gpt2",  # ~500MB - fallback if more memory available
+    "gpt2",  # ~500MB
+    "distilgpt2",  # ~250MB
+    "microsoft/DialoGPT-small",  # ~250MB
+    "huggingface/CodeBERTa-small-v1",  # Code tasks
 ]
 
-MIN_MEMORY_FOR_MODEL = 500  # Minimum MB required to attempt model loading
-MAX_TEST_CASES = 15  # Maximum number of test cases to return
-
-@dataclass
-class TestCase:
-    id: str
-    title: str
-    description: str
-    preconditions: List[str]
-    steps: List[str]
-    expected: str
-    postconditions: List[str]
-    test_data: str
-    priority: str = "Medium"
-    category: str = "Functional"
-
-# Enhanced pattern matching with more categories
-REQUIREMENT_PATTERNS = {
-    'authentication': {
-        'keywords': ['login', 'authentication', 'signin', 'sign in', 'password', 'username', 'credential', 'auth'],
-        'priority': 'High',
-        'category': 'Security',
-        'generator': 'generate_security_tests'
-    },
-    'authorization': {
-        'keywords': ['permission', 'role', 'access', 'privilege', 'authorize', 'admin', 'user level'],
-        'priority': 'High', 
-        'category': 'Security',
-        'generator': 'generate_security_tests'
-    },
-    'data_validation': {
-        'keywords': ['validate', 'validation', 'input', 'format', 'check', 'verify', 'constraint'],
-        'priority': 'High',
-        'category': 'Functional',
-        'generator': 'generate_validation_tests'
-    },
-    'database': {
-        'keywords': ['database', 'db', 'store', 'save', 'persist', 'record', 'data storage', 'crud'],
-        'priority': 'Medium',
-        'category': 'Data',
-        'generator': 'generate_data_tests'
-    },
-    'performance': {
-        'keywords': ['performance', 'speed', 'time', 'response', 'load', 'concurrent', 'scalability'],
-        'priority': 'Medium',
-        'category': 'Performance',
-        'generator': 'generate_performance_tests'
-    },
-    'api': {
-        'keywords': ['api', 'endpoint', 'service', 'request', 'response', 'rest', 'http'],
-        'priority': 'High',
-        'category': 'Integration',
-        'generator': 'generate_api_tests'
-    },
-    'error_handling': {
-        'keywords': ['error', 'exception', 'failure', 'invalid', 'incorrect', 'wrong'],
-        'priority': 'High',
-        'category': 'Reliability',
-        'generator': 'generate_error_tests'
-    },
-    'security': {
-        'keywords': ['security', 'encrypt', 'secure', 'ssl', 'https', 'token', 'session'],
-        'priority': 'High',
-        'category': 'Security',
-        'generator': 'generate_security_tests'
-    }
-}
-
-class TestCaseGenerator:
-    """Main class for generating test cases with AI and template fallback"""
-    
-    def __init__(self):
-        self.model_name = None
-        self.tokenizer = None
-        self.model = None
-        self._initialize_model()
-        
-    def _initialize_model(self):
-        """Initialize the optimal model based on available memory"""
-        available_mem = psutil.virtual_memory().available / (1024 * 1024)
-        logger.info(f"Available memory: {available_mem:.1f}MB")
-        
-        if available_mem < MIN_MEMORY_FOR_MODEL:
-            logger.warning("Insufficient memory for model loading, using template fallback")
-            return
-            
-        # Try models in order of preference
-        for model_name in MEMORY_OPTIMIZED_MODELS:
-            try:
-                self.tokenizer, self.model = self._load_model_safely(model_name)
-                if self.model:
-                    self.model_name = model_name
-                    logger.info(f"Successfully loaded model: {model_name}")
-                    break
-            except Exception as e:
-                logger.warning(f"Failed to load {model_name}: {str(e)}")
-                continue
-                
-    def _load_model_safely(self, model_name: str) -> Tuple[Optional[AutoTokenizer], Optional[AutoModelForCausalLM]]:
-        """Safely load model with memory optimizations"""
-        try:
-            logger.info(f"Attempting to load {model_name}")
-            
-            # Load tokenizer first
-            tokenizer = AutoTokenizer.from_pretrained(
-                model_name,
-                padding_side='left',
-                use_fast=True
-            )
-            
-            # Ensure pad token is set
-            if tokenizer.pad_token is None:
-                tokenizer.pad_token = tokenizer.eos_token if tokenizer.eos_token else '[PAD]'
-            
-            # Load model with optimized settings
-            model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-                low_cpu_mem_usage=True,
-                device_map="auto" if torch.cuda.is_available() else None
-            )
-            
-            # Explicitly move to CPU if needed
-            if not torch.cuda.is_available():
-                model = model.to('cpu')
-                
-            model.eval()
-            return tokenizer, model
-            
-        except Exception as e:
-            logger.error(f"Error loading model {model_name}: {str(e)}")
-            # Clean up if partial load occurred
-            if 'tokenizer' in locals():
-                del tokenizer
-            if 'model' in locals() and model:
-                del model
-            gc.collect()
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-            return None, None
-            
-    def generate_test_cases(self, srs_text: str) -> List[TestCase]:
-        """Generate test cases using best available method"""
-        # First try AI generation if model is available
-        if self.model and self.tokenizer:
-            try:
-                ai_cases = self._generate_with_ai(srs_text)
-                if ai_cases:
-                    logger.info("Successfully generated test cases with AI")
-                    return ai_cases[:MAX_TEST_CASES]
-            except Exception as e:
-                logger.warning(f"AI generation failed: {str(e)}, falling back to templates")
-                
-        # Fall back to template-based generation
-        return self._generate_with_templates(srs_text)[:MAX_TEST_CASES]
-        
-    def _generate_with_ai(self, srs_text: str) -> List[TestCase]:
-        """Generate test cases using AI model"""
-        max_input_length = 500  # Increased from 300 for better context
-        prompt = f"""Generate comprehensive test cases for these software requirements:
-{self._truncate_text(srs_text, max_input_length)}
-
-Provide test cases in this format:
-1. [Test Case Title]
-   - Description: [description]
-   - Steps: [step1; step2; step3]
-   - Expected: [expected result]
-   
-2. [Next Test Case Title]..."""
-        
-        try:
-            inputs = self.tokenizer(
-                prompt,
-                return_tensors="pt",
-                max_length=512,
-                truncation=True,
-                padding=True,
-                return_attention_mask=True
-            )
-            
-            # Generate with more controlled parameters
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    input_ids=inputs['input_ids'],
-                    attention_mask=inputs['attention_mask'],
-                    max_new_tokens=300,
-                    num_return_sequences=1,
-                    temperature=0.7,
-                    top_p=0.9,
-                    do_sample=True,
-                    pad_token_id=self.tokenizer.pad_token_id,
-                    eos_token_id=self.tokenizer.eos_token_id
-                )
-                
-            generated = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            return self._parse_ai_output(generated)
-            
-        except Exception as e:
-            logger.error(f"AI generation error: {str(e)}")
-            raise
-        finally:
-            # Clean up
-            if 'inputs' in locals():
-                del inputs
-            if 'outputs' in locals():
-                del outputs
-            gc.collect()
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-                
-    def _parse_ai_output(self, text: str) -> List[TestCase]:
-        """Parse AI-generated text into structured test cases"""
-        cases = []
-        current_case = None
-        
-        for line in text.split('\n'):
-            line = line.strip()
-            if line.startswith(('1.', '2.', '3.', '4.', '5.', '6.', '7.', '8.', '9.')):
-                if current_case:
-                    cases.append(current_case)
-                title = line[2:].strip()
-                current_case = TestCase(
-                    id=f"TC_AI_{len(cases)+1:03d}",
-                    title=title,
-                    description="",
-                    preconditions=["System is accessible"],
-                    steps=[],
-                    expected="",
-                    postconditions=["Test executed"],
-                    test_data="As specified in requirements",
-                    priority="Medium",
-                    category="Functional"
-                )
-            elif line.lower().startswith('description:') and current_case:
-                current_case.description = line[12:].strip()
-            elif line.lower().startswith('steps:') and current_case:
-                steps = line[6:].strip().split(';')
-                current_case.steps = [s.strip() for s in steps if s.strip()]
-            elif line.lower().startswith('expected:') and current_case:
-                current_case.expected = line[9:].strip()
-                
-        if current_case:
-            cases.append(current_case)
-            
-        return cases or [self._create_fallback_case()]
-        
-    def _generate_with_templates(self, srs_text: str) -> List[TestCase]:
-        """Generate test cases using pattern matching and templates"""
-        patterns = self._analyze_requirements(srs_text)
-        test_cases = []
-        
-        for pattern_name, pattern_data in patterns.items():
-            generator_name = REQUIREMENT_PATTERNS[pattern_name]['generator']
-            generator = getattr(self, generator_name, self._generate_generic_tests)
-            cases = generator(pattern_data['matches'])
-            
-            for i, case in enumerate(cases):
-                case.id = f"TC_{pattern_name.upper()}_{i+1:03d}"
-                case.priority = pattern_data['priority']
-                case.category = pattern_data['category']
-                test_cases.append(case)
-                
-        return test_cases or [self._create_fallback_case()]
-        
-    def _analyze_requirements(self, text: str) -> Dict[str, Any]:
-        """Analyze text to detect requirement patterns"""
-        text_lower = text.lower()
-        detected = {}
-        
-        for name, info in REQUIREMENT_PATTERNS.items():
-            matches = []
-            for kw in info['keywords']:
-                if kw in text_lower:
-                    # Find context around keyword
-                    context = re.findall(rf'.{{0,50}}{re.escape(kw)}.{{0,50}}', text_lower)
-                    matches.extend(context[:3])  # Limit contexts
-                    
-            if matches:
-                detected[name] = {
-                    'matches': matches,
-                    'priority': info['priority'],
-                    'category': info['category']
-                }
-                
-        return detected
-        
-    def _create_fallback_case(self) -> TestCase:
-        """Create a generic fallback test case"""
-        return TestCase(
-            id="TC_GEN_001",
-            title="General Functionality Test",
-            description="Verify basic system functionality",
-            preconditions=["System is accessible"],
-            steps=["Execute core functionality"],
-            expected="System behaves as expected",
-            postconditions=["Test completed"],
-            test_data="Standard test data",
-            priority="Medium",
-            category="Functional"
+# 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
+    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
+
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16,
+            device_map="cpu",
+            low_cpu_mem_usage=True,
+            use_cache=False,
         )
-        
-    def _truncate_text(self, text: str, max_length: int) -> str:
-        """Safely truncate text to maximum length"""
-        return text[:max_length] + '...' if len(text) > max_length else text
-        
-    # Template generators for different test types
-    def generate_security_tests(self, matches: List[str]) -> List[TestCase]:
-        """Generate security-related test cases"""
-        return [
-            TestCase(
-                id="",
-                title="Authentication Validation",
-                description="Verify proper authentication mechanism",
-                preconditions=["System has authentication configured"],
-                steps=[
-                    "Attempt login with valid credentials",
-                    "Attempt login with invalid credentials",
-                    "Verify session handling"
-                ],
-                expected="Valid login succeeds, invalid fails, sessions secure",
-                postconditions=["Security verified"],
-                test_data="Test credentials",
-                priority="High",
-                category="Security"
-            ),
-            # Additional security test cases...
+
+        model.eval()
+        model.gradient_checkpointing_enable()
+        logger.info(f"✅ Model {model_name} loaded successfully")
+        return tokenizer, model
+
+    except Exception as e:
+        logger.error(f"❌ Failed to load model {model_name}: {e}")
+        return None, None
+
+def extract_keywords(text):
+    common_keywords = [
+        'login', 'authentication', 'user', 'password', 'database', 'data',
+        'interface', 'api', 'function', 'feature', 'requirement', 'system',
+        'input', 'output', 'validation', 'error', 'security', 'performance'
+    ]
+    words = re.findall(r'\b\w+\b', text.lower())
+    return [word for word in words if word in common_keywords]
+
+def generate_template_based_test_cases(srs_text):
+    keywords = extract_keywords(srs_text)
+    test_cases = []
+
+    if any(word in keywords for word in ['login', 'authentication', 'user', 'password']):
+        test_cases.extend([
+            {
+                "id": "TC_001",
+                "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",
+                "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"
+            }
+        ])
+
+    if any(word in keywords for word in ['database', 'data', 'store', 'save']):
+        test_cases.append({
+            "id": "TC_003",
+            "title": "Data Storage Test",
+            "description": "Test data storage functionality",
+            "steps": ["Enter data", "Save data", "Verify storage"],
+            "expected": "Data should be stored correctly"
+        })
+
+    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"
+            }
         ]
-        
-    def generate_validation_tests(self, matches: List[str]) -> List[TestCase]:
-        """Generate data validation test cases"""
-        return [
-            TestCase(
-                id="",
-                title="Input Validation",
-                description="Verify input validation rules",
-                preconditions=["System accepts user input"],
-                steps=[
-                    "Enter valid input",
-                    "Enter invalid input",
-                    "Verify system response"
-                ],
-                expected="Valid input accepted, invalid rejected with messages",
-                postconditions=["Validation rules verified"],
-                test_data="Valid and invalid test data",
-                priority="High",
-                category="Functional"
+
+    return test_cases
+
+def parse_generated_test_cases(generated_text):
+    lines = generated_text.split('\n')
+    test_cases = []
+    current_case = {}
+    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 = {
+                "id": f"TC_{case_counter:03d}",
+                "title": line,
+                "description": line,
+                "steps": ["Execute the test"],
+                "expected": "Test should pass"
+            }
+            case_counter += 1
+
+    if current_case:
+        test_cases.append(current_case)
+
+    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"
+        }]
+
+    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:
+{srs_text}
+
+Test Cases:
+1."""
+
+    try:
+        inputs = tokenizer.encode(
+            prompt,
+            return_tensors="pt",
+            max_length=150,
+            truncation=True
+        )
+
+        with torch.no_grad():
+            outputs = model.generate(
+                inputs,
+                max_new_tokens=100,
+                num_return_sequences=1,
+                temperature=0.7,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id,
+                use_cache=False,
             )
-        ]
-        
-    # Additional generator methods for other test types...
-    # generate_performance_tests, generate_api_tests, etc.
 
-# Singleton instance
-_generator_instance = None
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        del inputs, outputs
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None
+        return parse_generated_test_cases(generated_text)
 
-def get_generator() -> TestCaseGenerator:
-    """Get the singleton generator instance"""
+    except Exception as e:
+        logger.error(f"❌ AI generation failed: {e}")
+        raise
+
+def generate_with_fallback(srs_text):
+    model_name = get_optimal_model_for_memory()
+
+    if model_name:
+        tokenizer, model = load_model_with_memory_optimization(model_name)
+        if tokenizer and model:
+            try:
+                test_cases = generate_with_ai_model(srs_text, tokenizer, model)
+                reason = get_algorithm_reason(model_name)
+                return test_cases, model_name, "transformer (causal LM)", reason
+            except Exception as e:
+                logger.warning(f"AI generation failed: {e}, falling back to templates")
+
+    logger.info("⚠️ Using fallback template-based generation")
+    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 get_generator():
     global _generator_instance
     if _generator_instance is None:
-        _generator_instance = TestCaseGenerator()
+        class Generator:
+            def __init__(self):
+                self.model_name = get_optimal_model_for_memory()
+                self.tokenizer = None
+                self.model = 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",
+                    "status": "loaded" if self.model else "template_mode",
+                    "memory_usage": f"{mem:.1f}MB",
+                    "optimization": "low_memory"
+                }
+
+        _generator_instance = Generator()
+
     return _generator_instance
 
-def generate_test_cases(srs_text: str) -> List[Dict]:
-    """Main API function to generate test cases"""
-    generator = get_generator()
-    cases = generator.generate_test_cases(srs_text)
-    return [case.__dict__ for case in cases]
-
-def generate_test_cases_and_info(srs_text: str) -> Dict[str, Any]:
-    """Generate test cases with metadata about generation method"""
-    generator = get_generator()
-    cases = generator.generate_test_cases(srs_text)
-    
+def monitor_memory():
+    mem = psutil.Process().memory_info().rss / 1024 / 1024
+    logger.info(f"Memory usage: {mem:.1f}MB")
+    if mem > 450:
+        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": generator.model_name or "Template-Based",
-        "algorithm": "AI" if generator.model else "Template",
-        "test_cases": [case.__dict__ for case in cases],
-        "memory_usage": f"{psutil.Process().memory_info().rss / (1024 * 1024):.1f}MB"
+        "model": model_name,
+        "algorithm": algorithm_used,
+        "reason": reason,
+        "test_cases": test_cases
     }
 
-# Example usage
-if __name__ == "__main__":
-    sample_reqs = """
-    The system must implement secure user authentication with password hashing.
-    All API endpoints must validate input data and return appropriate error codes.
-    The application should handle 100 concurrent users with response times under 2 seconds.
-    """
-    
-    print("Generating test cases...")
-    test_cases = generate_test_cases(sample_reqs)
-    for case in test_cases:
-        print(f"\n{case['id']}: {case['title']}")
-        print(f"Priority: {case['priority']}, Category: {case['category']}")
-        print(f"Steps: {case['steps']}")
+# ✅ 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."
+    elif model_name == "gpt2":
+        return "Chosen for general-purpose text generation with moderate memory headroom."
+    elif model_name is None:
+        return "No model used due to insufficient memory. Rule-based template generation chosen instead."
+    else:
+        return "Model selected based on best tradeoff between memory usage and language generation capability."