Update src/summarizer.py

src/summarizer.py

@@ -1,28 +1,28 @@
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqSummarization
 import torch
 import logging
 
 class TextSummarizer:
     def __init__(self, model_name="facebook/bart-large-cnn"):
         """
-        Initialize summarization pipeline
+        Initialize summarization model directly without using pipeline
         
         Args:
             model_name (str): Hugging Face model for summarization
         """
         try:
-            # Configure device
-            device = 0 if torch.cuda.is_available() else -1
-            logging.info(f"Using device: {'cuda' if device == 0 else 'cpu'}")
+            # Force CPU usage and disable GPU
+            self.device = torch.device('cpu')
             
-            # Initialize pipeline with explicit device mapping and lower precision
-            self.summarizer = pipeline(
-                "summarization", 
-                model=model_name,
-                device=device,
-                torch_dtype=torch.float32
-            )
-            logging.info("Summarization pipeline initialized successfully")
+            # Initialize tokenizer and model separately
+            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+            self.model = AutoModelForSeq2SeqSummarization.from_pretrained(model_name)
+            
+            # Move model to CPU and eval mode
+            self.model = self.model.to(self.device)
+            self.model.eval()
+            
+            logging.info("Summarization model initialized successfully")
             
         except Exception as e:
             logging.error(f"Failed to load summarization model: {str(e)}")
@@ -48,20 +48,33 @@ class TextSummarizer:
             # Ensure min_length is less than max_length
             min_length = min(min_length, max_length)
             
-            # Generate summary with chunking for long texts
+            # Process text in chunks due to length limitations
             max_chunk_length = 1024  # BART's max input length
             chunks = [text[i:i + max_chunk_length] for i in range(0, len(text), max_chunk_length)]
             summaries = []
             
-            for chunk in chunks:
-                if chunk.strip():
-                    summary = self.summarizer(
-                        chunk,
-                        max_length=max_length // len(chunks),  # Distribute length across chunks
-                        min_length=min_length // len(chunks),
-                        do_sample=False
-                    )[0]['summary_text']
-                    summaries.append(summary)
+            with torch.no_grad():  # Disable gradient calculation
+                for chunk in chunks:
+                    if chunk.strip():
+                        # Tokenize
+                        inputs = self.tokenizer(chunk, max_length=1024, truncation=True, 
+                                             return_tensors="pt")
+                        inputs = inputs.to(self.device)
+                        
+                        # Generate summary
+                        summary_ids = self.model.generate(
+                            inputs["input_ids"],
+                            num_beams=4,
+                            max_length=max_length // len(chunks),
+                            min_length=min_length // len(chunks),
+                            length_penalty=2.0,
+                            early_stopping=True
+                        )
+                        
+                        # Decode summary
+                        summary = self.tokenizer.decode(summary_ids[0], 
+                                                      skip_special_tokens=True)
+                        summaries.append(summary)
             
             return " ".join(summaries)