Quintino Fernandes commited on
Commit
e773b70
·
1 Parent(s): 94c2660
Files changed (1) hide show
  1. models/summarization.py +1 -1
models/summarization.py CHANGED
@@ -4,7 +4,7 @@ import torch
4
  class SummarizationModel:
5
  def __init__(self):
6
  self.device = "cuda" if torch.cuda.is_available() else "cpu"
7
- self.tokenizer = T5Tokenizer.from_pretrained('unicamp-dl/ptt5-base-portuguese-vocab').to(self.device)
8
  self.model = T5ForConditionalGeneration.from_pretrained('recogna-nlp/ptt5-base-summ').to(self.device)
9
 
10
  def summarize(self, text: str) -> str:
 
4
  class SummarizationModel:
5
  def __init__(self):
6
  self.device = "cuda" if torch.cuda.is_available() else "cpu"
7
+ self.tokenizer = T5Tokenizer.from_pretrained('unicamp-dl/ptt5-base-portuguese-vocab')
8
  self.model = T5ForConditionalGeneration.from_pretrained('recogna-nlp/ptt5-base-summ').to(self.device)
9
 
10
  def summarize(self, text: str) -> str: