Update src/models/summarization.py

src/models/summarization.py

@@ -16,12 +16,11 @@ class Summarizer:
     def load_model(self):
         """Load the fine-tuned BART summarization model."""
         try:
+            with open('bart_ami_finetuned.pkl','rb') as f:
+                self.model = pickle.load(f)
             # Load the tokenizer
             self.tokenizer = BartTokenizer.from_pretrained("facebook/bart-base")
             
-            # Load the fine-tuned model
-            self.model = BartForConditionalGeneration.from_pretrained("bart_ami_finetuned.pkl")
-            
             # Move model to appropriate device (GPU if available)
             self.model.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
             return self.model