Increased Model Efficiency

model/analyzer.py

@@ -1,13 +1,13 @@
 import os
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import asyncio
 import torch
 from datetime import datetime
 import gradio as gr
 from typing import Dict, List, Union, Optional
 import logging
 import traceback
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
-# Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -19,63 +19,62 @@ class ContentAnalyzer:
         logger.info(f"Initialized analyzer with device: {self.device}")
 
     async def load_model(self, progress=None) -> None:
-        """Load the model and tokenizer with progress updates and detailed logging."""
+        """Load quantized model with optimized configuration."""
         try:
-            print("\n=== Starting Model Loading ===")
-            print(f"Time: {datetime.now()}")
-            
             if progress:
                 progress(0.1, "Loading tokenizer...")
             
-            print("Loading tokenizer...")
+            # Quantization configuration
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=torch.float16,
+                bnb_4bit_quant_type="nf4"
+            )
+
             self.tokenizer = AutoTokenizer.from_pretrained(
                 "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
                 use_fast=True
             )
 
             if progress:
-                progress(0.3, "Loading model...")
-            
-            print(f"Loading model on {self.device}...")
+                progress(0.3, "Loading quantized model...")
+
             self.model = AutoModelForCausalLM.from_pretrained(
                 "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
-                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
+                quantization_config=quantization_config,
                 device_map="auto"
             )
 
             if progress:
                 progress(0.5, "Model loaded successfully")
-            
-            print("Model and tokenizer loaded successfully")
-            logger.info(f"Model loaded successfully on {self.device}")
+
         except Exception as e:
-            logger.error(f"Error loading model: {str(e)}")
-            print(f"\nERROR DURING MODEL LOADING: {str(e)}")
-            print("Stack trace:")
+            logger.error(f"Model loading error: {str(e)}")
             traceback.print_exc()
             raise
 
-    def _chunk_text(self, text: str, chunk_size: int = 2048, overlap: int = 256) -> List[str]:
-        """Split text into overlapping chunks for processing."""
+    def _semantic_chunk_text(self, text: str, max_chunk_size: int = 4096) -> List[str]:
+        """Semantic chunking with dynamic sizing."""
         chunks = []
-        for i in range(0, len(text), chunk_size - overlap):
-            chunk = text[i:i + chunk_size]
-            chunks.append(chunk)
-        print(f"Split text into {len(chunks)} chunks with {overlap} token overlap")
+        current_chunk = ""
+        for sentence in text.split('.'):
+            if len(current_chunk) + len(sentence) < max_chunk_size:
+                current_chunk += sentence + '.'
+            else:
+                chunks.append(current_chunk.strip())
+                current_chunk = sentence + '.'
+        
+        if current_chunk:
+            chunks.append(current_chunk.strip())
+        
         return chunks
 
     async def analyze_chunk(
-        self,
-        chunk: str,
-        progress: Optional[gr.Progress] = None,
-        current_progress: float = 0,
-        progress_step: float = 0
+        self, 
+        chunk: str, 
+        progress: Optional[gr.Progress] = None
     ) -> List[str]:
-        """Analyze a single chunk of text for triggers with detailed logging."""
-        print(f"\n--- Processing Chunk ---")
-        print(f"Chunk text (preview): {chunk[:50]}...")
-        
-        # Comprehensive trigger categories
+        """Optimized single-pass chunk analysis."""
         categories = [
             "Violence", "Death", "Substance Use", "Gore", 
             "Vomit", "Sexual Content", "Sexual Abuse", 
@@ -83,111 +82,67 @@ class ContentAnalyzer:
             "Mental Health Issues"
         ]
 
-        # Comprehensive prompt for single-pass analysis
-        prompt = f"""Comprehensive Content Sensitivity Analysis
-
-Carefully analyze the following text for sensitive content categories:
-{', '.join(categories)}
-
-Detailed Requirements:
-1. Thoroughly examine entire text chunk
-2. Identify presence of ANY of these categories
-3. Provide clear, objective assessment
-4. Minimal subjective interpretation
-
-TEXT CHUNK:
-{chunk}
-
-RESPONSE FORMAT:
-- List categories DEFINITIVELY present
-- Brief objective justification for each
-- Strict YES/NO categorization"""
+        prompt = f"""Analyze this text for sensitive content.
+Categories: {', '.join(categories)}
+Identify ALL present categories.
+Be precise and direct.
+Chunk: {chunk}
+Output Format: Comma-separated category names if present."""
 
         try:
-            print("Sending prompt to model...")
-            inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=4096)
+            inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True)
             inputs = {k: v.to(self.device) for k, v in inputs.items()}
 
-            with torch.no_grad():
-                print("Generating response...")
-                outputs = self.model.generate(
-                    **inputs,
-                    max_new_tokens=256,
-                    do_sample=True,
-                    temperature=0.2,
-                    top_p=0.9,
-                    pad_token_id=self.tokenizer.eos_token_id
-                )
-
-            response_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
-            print("Full Model Response:", response_text)
-
-            # Parse detected triggers
-            detected_triggers = []
-            for category in categories:
-                if category.upper() in response_text.upper():
-                    detected_triggers.append(category)
-
-            print(f"Detected triggers in chunk: {detected_triggers}")
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=128,
+                do_sample=True,
+                temperature=0.2,
+                top_p=0.9,
+                pad_token_id=self.tokenizer.eos_token_id
+            )
 
-            if progress:
-                current_progress += progress_step
-                progress(min(current_progress, 0.9), "Analyzing chunk...")
+            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+            # Extract detected categories
+            detected = [
+                cat for cat in categories 
+                if cat.upper() in response.upper()
+            ]
 
-            return detected_triggers
+            return detected
 
         except Exception as e:
-            logger.error(f"Error analyzing chunk: {str(e)}")
-            print(f"Error during chunk analysis: {str(e)}")
-            traceback.print_exc()
+            logger.error(f"Chunk analysis error: {str(e)}")
             return []
 
     async def analyze_script(self, script: str, progress: Optional[gr.Progress] = None) -> List[str]:
-        """Analyze the entire script for triggers with progress updates."""
-        print("\n=== Starting Script Analysis ===")
-        print(f"Time: {datetime.now()}")
-
         if not self.model or not self.tokenizer:
             await self.load_model(progress)
 
-        chunks = self._chunk_text(script)
-        identified_triggers = set()
-        progress_step = 0.4 / len(chunks)
-        current_progress = 0.5  # Starting after model loading
-
-        for chunk_idx, chunk in enumerate(chunks, 1):
-            chunk_triggers = await self.analyze_chunk(
-                chunk,
-                progress,
-                current_progress,
-                progress_step
-            )
-            identified_triggers.update(chunk_triggers)
-
-        if progress:
-            progress(0.95, "Finalizing results...")
+        chunks = self._semantic_chunk_text(script)
+        
+        # Concurrent chunk processing
+        tasks = [self.analyze_chunk(chunk) for chunk in chunks]
+        chunk_results = await asyncio.gather(*tasks)
 
-        final_triggers = list(identified_triggers)
-        print("\n=== Analysis Complete ===")
-        print("Final Results:", final_triggers)
+        # Flatten and deduplicate results
+        identified_triggers = set(
+            trigger 
+            for chunk_triggers in chunk_results 
+            for trigger in chunk_triggers
+        )
 
-        return final_triggers if final_triggers else ["None"]
+        return list(identified_triggers) or ["None"]
 
 async def analyze_content(
     script: str,
     progress: Optional[gr.Progress] = None
 ) -> Dict[str, Union[List[str], str]]:
-    """Main analysis function for the Gradio interface."""
-    print("\n=== Starting Content Analysis ===")
-    print(f"Time: {datetime.now()}")
-    
     analyzer = ContentAnalyzer()
     
     try:
         triggers = await analyzer.analyze_script(script, progress)
-        
-        if progress:
-            progress(1.0, "Analysis complete!")
 
         result = {
             "detected_triggers": triggers,
@@ -196,14 +151,10 @@ async def analyze_content(
             "analysis_timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
         }
 
-        print("\nFinal Result Dictionary:", result)
         return result
 
     except Exception as e:
         logger.error(f"Analysis error: {str(e)}")
-        print(f"\nERROR OCCURRED: {str(e)}")
-        print("Stack trace:")
-        traceback.print_exc()
         return {
             "detected_triggers": ["Error occurred during analysis"],
             "confidence": "Error",
@@ -213,7 +164,6 @@ async def analyze_content(
         }
 
 if __name__ == "__main__":
-    # Gradio interface
     iface = gr.Interface(
         fn=analyze_content,
         inputs=gr.Textbox(lines=8, label="Input Text"),