Model changes

models/summarization.py

@@ -4,10 +4,12 @@ import torch
 class SummarizationModel:
     def __init__(self):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.tokenizer = T5Tokenizer.from_pretrained('unicamp-dl/ptt5-base-portuguese-vocab')
+        self.tokenizer = T5Tokenizer.from_pretrained('unicamp-dl/ptt5-base-portuguese-vocab').to(self.device)
         self.model = T5ForConditionalGeneration.from_pretrained('recogna-nlp/ptt5-base-summ').to(self.device)
     
-    def summarize(self, text: str, max_length: int = 256, min_length: int = 128) -> str:
+    def summarize(self, text: str) -> str:
+        print(next(self.tokenizer.parameters()).device)
+        print(next(self.model.parameters()).device)
         """Summarize the input text using T5 model"""
         # Model and tokenization parameters
         inputs = self.tokenizer.encode(
@@ -19,8 +21,8 @@ class SummarizationModel:
         
         summary_ids = self.model.generate(
             inputs,
-            max_length=max_length,
-            min_length=min_length,
+            max_length=256,
+            min_length=64,
             num_beams=4,
             no_repeat_ngram_size=3,
             early_stopping=True,