changed the bart model from large to base

config/config.py

@@ -34,6 +34,18 @@ class Settings:
         # Better for development purposes.
         return os.getenv('EMBEDDING_MODEL', 'all-MiniLM-L6-v2')
 
+    # New Conversation Summarizer Settings
+    SUMMARIZER_CONFIG = {
+        # 'facebook/bart-large-cnn', for bigger and better model
+        'model_name': os.getenv('SUMMARIZER_MODEL', 'facebook/bart-base'),
+        'max_length': int(os.getenv('SUMMARIZER_MAX_LENGTH', '130')),
+        'min_length': int(os.getenv('SUMMARIZER_MIN_LENGTH', '30')),
+        'device': -1,  # CPU
+        'model_kwargs': {
+            'low_cpu_mem_usage': True
+        }
+    }
+
     # Vector Store Configuration
     CHROMA_PATH = os.getenv('CHROMA_PATH', './chroma_db')
 

src/utils/conversation_summarizer.py

@@ -3,22 +3,36 @@ from typing import List, Dict
 from transformers import pipeline
 import numpy as np
 from datetime import datetime
+from config.config import settings
+
 
 class ConversationSummarizer:
     def __init__(
         self,
-        model_name: str = "facebook/bart-large-cnn",
-        max_length: int = 130,
-        min_length: int = 30
+        model_name: str = None,
+        max_length: int = None,
+        min_length: int = None
     ):
-        """Initialize the summarizer"""
+        """
+        Initialize the summarizer
+
+        Args:
+            model_name (str, optional): Override default model from config
+            max_length (int, optional): Override default max_length from config
+            min_length (int, optional): Override default min_length from config
+        """
+        # Use provided values or fall back to config values
+        self.model_name = model_name or settings.SUMMARIZER_CONFIG['model_name']
+        self.max_length = max_length or settings.SUMMARIZER_CONFIG['max_length']
+        self.min_length = min_length or settings.SUMMARIZER_CONFIG['min_length']
+
+        # Initialize the summarizer with config settings
         self.summarizer = pipeline(
             "summarization",
-            model=model_name,
-            device=-1  # CPU
+            model=self.model_name,
+            device=settings.SUMMARIZER_CONFIG['device'],
+            model_kwargs=settings.SUMMARIZER_CONFIG['model_kwargs']
         )
-        self.max_length = max_length
-        self.min_length = min_length
 
     async def summarize_conversation(
         self,
@@ -30,7 +44,7 @@ class ConversationSummarizer:
         """
         # Format conversation for summarization
         formatted_convo = self._format_conversation(messages)
-        
+
         # Generate summary
         summary = self.summarizer(
             formatted_convo,
@@ -38,13 +52,14 @@ class ConversationSummarizer:
             min_length=self.min_length,
             do_sample=False
         )[0]['summary_text']
-        
+
         # Extract key insights
         insights = self._extract_insights(messages)
-        
+
         # Generate metadata if requested
-        metadata = self._generate_metadata(messages) if include_metadata else {}
-        
+        metadata = self._generate_metadata(
+            messages) if include_metadata else {}
+
         return {
             'summary': summary,
             'key_insights': insights,
@@ -58,7 +73,7 @@ class ConversationSummarizer:
             role = msg.get('role', 'unknown')
             content = msg.get('content', '')
             formatted.append(f"{role}: {content}")
-        
+
         return "\n".join(formatted)
 
     def _extract_insights(self, messages: List[Dict]) -> Dict:
@@ -68,13 +83,13 @@ class ConversationSummarizer:
             'user': len([m for m in messages if m.get('role') == 'user']),
             'assistant': len([m for m in messages if m.get('role') == 'assistant'])
         }
-        
+
         # Calculate average message length
         avg_length = np.mean([len(m.get('content', '')) for m in messages])
-        
+
         # Extract main topics (simplified)
         topics = self._extract_topics(messages)
-        
+
         return {
             'message_distribution': message_counts,
             'average_message_length': int(avg_length),
@@ -86,7 +101,7 @@ class ConversationSummarizer:
         """Extract main topics from conversation"""
         # Combine all messages
         full_text = " ".join([m.get('content', '') for m in messages])
-        
+
         # Use the summarizer to extract main points
         topics = self.summarizer(
             full_text,
@@ -94,14 +109,14 @@ class ConversationSummarizer:
             min_length=10,
             do_sample=False
         )[0]['summary_text'].split('. ')
-        
+
         return topics
 
     def _generate_metadata(self, messages: List[Dict]) -> Dict:
         """Generate conversation metadata"""
         if not messages:
             return {}
-            
+
         return {
             'start_time': messages[0].get('timestamp', None),
             'end_time': messages[-1].get('timestamp', None),
@@ -112,8 +127,10 @@ class ConversationSummarizer:
     def _calculate_duration(self, messages: List[Dict]) -> float:
         """Calculate conversation duration in minutes"""
         try:
-            start_time = datetime.fromisoformat(messages[0].get('timestamp', ''))
-            end_time = datetime.fromisoformat(messages[-1].get('timestamp', ''))
+            start_time = datetime.fromisoformat(
+                messages[0].get('timestamp', ''))
+            end_time = datetime.fromisoformat(
+                messages[-1].get('timestamp', ''))
             return (end_time - start_time).total_seconds() / 60
         except:
             return 0
@@ -125,4 +142,4 @@ class ConversationSummarizer:
             if message.get('sources'):
                 for source in message['sources']:
                     sources.add(source.get('filename', ''))
-        return list(sources)
+        return list(sources)