Update models/summary_models.py

models/summary_models.py

@@ -1,249 +1,270 @@
-# models/summary_models.py
-import logging
-from typing import Dict, List, Optional, Tuple, Union, Any
-import torch
-from transformers import T5Tokenizer, T5ForConditionalGeneration
-
-class SummaryModelManager:
-    def __init__(self, token_manager=None, cache_manager=None, metrics_calculator=None):
-        """Initialize the SummaryModelManager with optional utilities."""
-        self.logger = logging.getLogger(__name__)
-        self.token_manager = token_manager
-        self.cache_manager = cache_manager
-        self.metrics_calculator = metrics_calculator
-        
-        # Model instance
-        self.model = None
-        self.tokenizer = None
-        
-        # Model name
-        self.model_name = "t5-small"
-        
-        # Track initialization state
-        self.initialized = False
-        
-        # Default generation parameters
-        self.default_params = {
-            "max_length": 150,
-            "min_length": 40,
-            "length_penalty": 2.0,
-            "num_beams": 4,
-            "early_stopping": True
-        }
-        
-    def initialize_model(self):
-        """Initialize the summarization model."""
-        if self.initialized:
-            return
-            
-        try:
-            # Register with token manager if available
-            if self.token_manager:
-                self.token_manager.register_model(
-                    self.model_name, "summarization")
-                
-            # Load model and tokenizer
-            self.logger.info(f"Loading summary model: {self.model_name}")
-            self.tokenizer = T5Tokenizer.from_pretrained(self.model_name)
-            self.model = T5ForConditionalGeneration.from_pretrained(self.model_name)
-                
-            self.initialized = True
-            self.logger.info("Summary model initialized successfully")
-            
-        except Exception as e:
-            self.logger.error(f"Failed to initialize summary model: {e}")
-            raise
-    
-    def generate_summary(self, text: str, prefix: str = "summarize: ", 
-                        agent_name: str = "report_generation", 
-                        params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
-        """
-        Generate a summary of the given text.
-        Returns the summary and metadata.
-        """
-        # Initialize model if needed
-        if not self.initialized:
-            self.initialize_model()
-            
-        # Prepare input text
-        input_text = f"{prefix}{text}"
-        
-        # Check cache if available
-        if self.cache_manager:
-            cache_key = input_text[:100] + str(hash(input_text))  # Use prefix of text + hash as key
-            cache_hit, cached_result = self.cache_manager.get(
-                cache_key, namespace="summaries")
-                
-            if cache_hit:
-                # Update metrics if available
-                if self.metrics_calculator:
-                    self.metrics_calculator.update_cache_metrics(1, 0, 0.005)  # Estimated energy saving
-                return cached_result
-        
-        # Request token budget if available
-        if self.token_manager:
-            approved, reason = self.token_manager.request_tokens(
-                agent_name, "summarization", input_text, self.model_name)
-                
-            if not approved:
-                self.logger.warning(f"Token budget exceeded: {reason}")
-                return {"summary": "Token budget exceeded", "error": reason}
-        
-        # Tokenize
-        inputs = self.tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
-        
-        # Merge default and custom parameters
-        generation_params = self.default_params.copy()
-        if params:
-            generation_params.update(params)
-            
-        # Generate summary
-        with torch.no_grad():
-            output_ids = self.model.generate(
-                inputs.input_ids,
-                **generation_params
-            )
-            
-        # Decode summary
-        summary = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
-        
-        # Calculate compression ratio
-        input_length = len(text.split())
-        summary_length = len(summary.split())
-        compression_ratio = input_length / max(summary_length, 1)
-        
-        # Prepare result
-        result = {
-            "summary": summary,
-            "input_length": input_length,
-            "summary_length": summary_length,
-            "compression_ratio": compression_ratio
-        }
-        
-        # Log token usage if available
-        if self.token_manager:
-            input_tokens = len(inputs.input_ids[0])
-            output_tokens = len(output_ids[0])
-            total_tokens = input_tokens + output_tokens
-            
-            self.token_manager.log_usage(
-                agent_name, "summarization", total_tokens, self.model_name)
-                
-            # Log energy usage if metrics calculator is available
-            if self.metrics_calculator:
-                energy_usage = self.token_manager.calculate_energy_usage(
-                    total_tokens, self.model_name)
-                self.metrics_calculator.log_energy_usage(
-                    energy_usage, self.model_name, agent_name, "summarization")
-        
-        # Store in cache if available
-        if self.cache_manager:
-            self.cache_manager.put(cache_key, result, namespace="summaries")
-            
-        return result
-    
-    def generate_executive_summary(self, detailed_content: str, confidence_level: float,
-                                  agent_name: str = "report_generation") -> Dict[str, Any]:
-        """
-        Generate an executive summary with confidence indication.
-        Adjusts detail level based on confidence.
-        """
-        # Prepare prompt based on confidence
-        if confidence_level >= 0.7:
-            prefix = "summarize with high confidence: "
-            params = {"min_length": 30, "max_length": 100}
-        elif confidence_level >= 0.4:
-            prefix = "summarize with moderate confidence: "
-            params = {"min_length": 20, "max_length": 80}
-        else:
-            prefix = "summarize with low confidence: "
-            params = {"min_length": 15, "max_length": 60}
-            
-        # Generate summary
-        result = self.generate_summary(detailed_content, prefix=prefix, 
-                                      agent_name=agent_name, params=params)
-        
-        # Add confidence level to result
-        result["confidence_level"] = confidence_level
-        
-        # Add confidence statement
-        confidence_statement = self._generate_confidence_statement(confidence_level)
-        result["confidence_statement"] = confidence_statement
-        
-        return result
-    
-    def _generate_confidence_statement(self, confidence_level: float) -> str:
-        """Generate an appropriate confidence statement based on the level."""
-        if confidence_level >= 0.8:
-            return "This analysis is provided with high confidence based on strong evidence in the provided materials."
-        elif confidence_level >= 0.6:
-            return "This analysis is provided with good confidence based on substantial evidence in the provided materials."
-        elif confidence_level >= 0.4:
-            return "This analysis is provided with moderate confidence. Some aspects may require additional verification."
-        elif confidence_level >= 0.2:
-            return "This analysis is provided with limited confidence due to sparse relevant information in the provided materials."
-        else:
-            return "This analysis is provided with very low confidence due to insufficient relevant information in the provided materials."
-    
-    def combine_analyses(self, text_analyses: List[Dict[str, Any]], 
-                        image_analyses: List[Dict[str, Any]],
-                        topic: str, agent_name: str = "report_generation") -> Dict[str, Any]:
-        """
-        Combine text and image analyses into a coherent report.
-        Returns the combined report with metadata.
-        """
-        # Build combined content
-        combined_content = f"Topic: {topic}\n\n"
-        
-        # Add text analyses
-        combined_content += "Text Analysis:\n"
-        for i, analysis in enumerate(text_analyses):
-            if "error" in analysis:
-                continue
-            combined_content += f"- Document {i+1}: {analysis.get('summary', 'No summary available')}\n"
-            
-        # Add image analyses
-        combined_content += "\nImage Analysis:\n"
-        for i, analysis in enumerate(image_analyses):
-            if "error" in analysis:
-                continue
-            combined_content += f"- Image {i+1}: {analysis.get('caption', 'No caption available')}\n"
-            
-        # Calculate overall confidence based on analyses
-        text_confidence = sum(a.get("confidence", 0) for a in text_analyses) / max(len(text_analyses), 1)
-        image_confidence = sum(a.get("confidence", 0) for a in image_analyses) / max(len(image_analyses), 1)
-        
-        # Weight confidence (text analyses typically more important for deep dives)
-        overall_confidence = 0.7 * text_confidence + 0.3 * image_confidence
-        
-        # Generate detailed report
-        detailed_report = self.generate_summary(
-            combined_content, 
-            prefix=f"generate detailed report about {topic}: ", 
-            agent_name=agent_name,
-            params={"max_length": 300, "min_length": 100}
-        )
-        
-        # Generate executive summary
-        executive_summary = self.generate_executive_summary(
-            detailed_report["summary"], 
-            overall_confidence,
-            agent_name
-        )
-        
-        # Combine results
-        result = {
-            "topic": topic,
-            "executive_summary": executive_summary["summary"],
-            "confidence_statement": executive_summary["confidence_statement"],
-            "detailed_report": detailed_report["summary"],
-            "confidence_level": overall_confidence,
-            "text_confidence": text_confidence,
-            "image_confidence": image_confidence,
-            "source_count": {
-                "text": len(text_analyses),
-                "images": len(image_analyses)
-            }
-        }
-        
-        return result
+# models/summary_models.py
+import logging
+from typing import Dict, List, Optional, Tuple, Union, Any
+import torch
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+
+class SummaryModelManager:
+    def __init__(self, token_manager=None, cache_manager=None, metrics_calculator=None):
+        """Initialize the SummaryModelManager with optional utilities."""
+        self.logger = logging.getLogger(__name__)
+        self.token_manager = token_manager
+        self.cache_manager = cache_manager
+        self.metrics_calculator = metrics_calculator
+        
+        # Model instance
+        self.model = None
+        self.tokenizer = None
+        
+        # Model name
+        self.model_name = "t5-small"
+        
+        # Track initialization state
+        self.initialized = False
+        
+        # Default generation parameters
+        self.default_params = {
+            "max_length": 150,
+            "min_length": 40,
+            "length_penalty": 2.0,
+            "num_beams": 4,
+            "early_stopping": True
+        }
+        
+    def initialize_model(self):
+        """Initialize the summarization model."""
+        if self.initialized:
+            return
+            
+        try:
+            # Register with token manager if available
+            if self.token_manager:
+                self.token_manager.register_model(
+                    self.model_name, "summarization")
+                
+            # Load model and tokenizer
+            self.logger.info(f"Loading summary model: {self.model_name}")
+            self.tokenizer = T5Tokenizer.from_pretrained(self.model_name)
+            self.model = T5ForConditionalGeneration.from_pretrained(self.model_name)
+                
+            self.initialized = True
+            self.logger.info("Summary model initialized successfully")
+            
+        except Exception as e:
+            #self.logger.error(f"Failed to initialize summary model: {e}")
+            #raise
+            # Try a fallback model that doesn't require SentencePiece
+            try:
+                fallback_model = "facebook/bart-base"
+                self.logger.info(f"Trying fallback model: {fallback_model}")
+    
+                from transformers import BartTokenizer, BartForConditionalGeneration
+    
+                self.tokenizer = BartTokenizer.from_pretrained(fallback_model)
+                self.model = BartForConditionalGeneration.from_pretrained(fallback_model)
+                self.model_name = fallback_model
+    
+                # Register fallback with token manager
+                if self.token_manager:
+                    self.token_manager.register_model(
+                        self.model_name, "summarization")
+    
+                self.initialized = True
+                self.logger.info("Fallback summary model initialized successfully")
+            except Exception as fallback_error:
+                self.logger.error(f"Failed to initialize fallback model: {fallback_error}")
+                raise
+    
+    def generate_summary(self, text: str, prefix: str = "summarize: ", 
+                        agent_name: str = "report_generation", 
+                        params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+        """
+        Generate a summary of the given text.
+        Returns the summary and metadata.
+        """
+        # Initialize model if needed
+        if not self.initialized:
+            self.initialize_model()
+            
+        # Prepare input text
+        input_text = f"{prefix}{text}"
+        
+        # Check cache if available
+        if self.cache_manager:
+            cache_key = input_text[:100] + str(hash(input_text))  # Use prefix of text + hash as key
+            cache_hit, cached_result = self.cache_manager.get(
+                cache_key, namespace="summaries")
+                
+            if cache_hit:
+                # Update metrics if available
+                if self.metrics_calculator:
+                    self.metrics_calculator.update_cache_metrics(1, 0, 0.005)  # Estimated energy saving
+                return cached_result
+        
+        # Request token budget if available
+        if self.token_manager:
+            approved, reason = self.token_manager.request_tokens(
+                agent_name, "summarization", input_text, self.model_name)
+                
+            if not approved:
+                self.logger.warning(f"Token budget exceeded: {reason}")
+                return {"summary": "Token budget exceeded", "error": reason}
+        
+        # Tokenize
+        inputs = self.tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
+        
+        # Merge default and custom parameters
+        generation_params = self.default_params.copy()
+        if params:
+            generation_params.update(params)
+            
+        # Generate summary
+        with torch.no_grad():
+            output_ids = self.model.generate(
+                inputs.input_ids,
+                **generation_params
+            )
+            
+        # Decode summary
+        summary = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        
+        # Calculate compression ratio
+        input_length = len(text.split())
+        summary_length = len(summary.split())
+        compression_ratio = input_length / max(summary_length, 1)
+        
+        # Prepare result
+        result = {
+            "summary": summary,
+            "input_length": input_length,
+            "summary_length": summary_length,
+            "compression_ratio": compression_ratio
+        }
+        
+        # Log token usage if available
+        if self.token_manager:
+            input_tokens = len(inputs.input_ids[0])
+            output_tokens = len(output_ids[0])
+            total_tokens = input_tokens + output_tokens
+            
+            self.token_manager.log_usage(
+                agent_name, "summarization", total_tokens, self.model_name)
+                
+            # Log energy usage if metrics calculator is available
+            if self.metrics_calculator:
+                energy_usage = self.token_manager.calculate_energy_usage(
+                    total_tokens, self.model_name)
+                self.metrics_calculator.log_energy_usage(
+                    energy_usage, self.model_name, agent_name, "summarization")
+        
+        # Store in cache if available
+        if self.cache_manager:
+            self.cache_manager.put(cache_key, result, namespace="summaries")
+            
+        return result
+    
+    def generate_executive_summary(self, detailed_content: str, confidence_level: float,
+                                  agent_name: str = "report_generation") -> Dict[str, Any]:
+        """
+        Generate an executive summary with confidence indication.
+        Adjusts detail level based on confidence.
+        """
+        # Prepare prompt based on confidence
+        if confidence_level >= 0.7:
+            prefix = "summarize with high confidence: "
+            params = {"min_length": 30, "max_length": 100}
+        elif confidence_level >= 0.4:
+            prefix = "summarize with moderate confidence: "
+            params = {"min_length": 20, "max_length": 80}
+        else:
+            prefix = "summarize with low confidence: "
+            params = {"min_length": 15, "max_length": 60}
+            
+        # Generate summary
+        result = self.generate_summary(detailed_content, prefix=prefix, 
+                                      agent_name=agent_name, params=params)
+        
+        # Add confidence level to result
+        result["confidence_level"] = confidence_level
+        
+        # Add confidence statement
+        confidence_statement = self._generate_confidence_statement(confidence_level)
+        result["confidence_statement"] = confidence_statement
+        
+        return result
+    
+    def _generate_confidence_statement(self, confidence_level: float) -> str:
+        """Generate an appropriate confidence statement based on the level."""
+        if confidence_level >= 0.8:
+            return "This analysis is provided with high confidence based on strong evidence in the provided materials."
+        elif confidence_level >= 0.6:
+            return "This analysis is provided with good confidence based on substantial evidence in the provided materials."
+        elif confidence_level >= 0.4:
+            return "This analysis is provided with moderate confidence. Some aspects may require additional verification."
+        elif confidence_level >= 0.2:
+            return "This analysis is provided with limited confidence due to sparse relevant information in the provided materials."
+        else:
+            return "This analysis is provided with very low confidence due to insufficient relevant information in the provided materials."
+    
+    def combine_analyses(self, text_analyses: List[Dict[str, Any]], 
+                        image_analyses: List[Dict[str, Any]],
+                        topic: str, agent_name: str = "report_generation") -> Dict[str, Any]:
+        """
+        Combine text and image analyses into a coherent report.
+        Returns the combined report with metadata.
+        """
+        # Build combined content
+        combined_content = f"Topic: {topic}\n\n"
+        
+        # Add text analyses
+        combined_content += "Text Analysis:\n"
+        for i, analysis in enumerate(text_analyses):
+            if "error" in analysis:
+                continue
+            combined_content += f"- Document {i+1}: {analysis.get('summary', 'No summary available')}\n"
+            
+        # Add image analyses
+        combined_content += "\nImage Analysis:\n"
+        for i, analysis in enumerate(image_analyses):
+            if "error" in analysis:
+                continue
+            combined_content += f"- Image {i+1}: {analysis.get('caption', 'No caption available')}\n"
+            
+        # Calculate overall confidence based on analyses
+        text_confidence = sum(a.get("confidence", 0) for a in text_analyses) / max(len(text_analyses), 1)
+        image_confidence = sum(a.get("confidence", 0) for a in image_analyses) / max(len(image_analyses), 1)
+        
+        # Weight confidence (text analyses typically more important for deep dives)
+        overall_confidence = 0.7 * text_confidence + 0.3 * image_confidence
+        
+        # Generate detailed report
+        detailed_report = self.generate_summary(
+            combined_content, 
+            prefix=f"generate detailed report about {topic}: ", 
+            agent_name=agent_name,
+            params={"max_length": 300, "min_length": 100}
+        )
+        
+        # Generate executive summary
+        executive_summary = self.generate_executive_summary(
+            detailed_report["summary"], 
+            overall_confidence,
+            agent_name
+        )
+        
+        # Combine results
+        result = {
+            "topic": topic,
+            "executive_summary": executive_summary["summary"],
+            "confidence_statement": executive_summary["confidence_statement"],
+            "detailed_report": detailed_report["summary"],
+            "confidence_level": overall_confidence,
+            "text_confidence": text_confidence,
+            "image_confidence": image_confidence,
+            "source_count": {
+                "text": len(text_analyses),
+                "images": len(image_analyses)
+            }
+        }
+        
+        return result