garvitcpp commited on
Commit
8199fc6
·
verified ·
1 Parent(s): f98eaa3

Update src/summarizer.py

Browse files
Files changed (1) hide show
  1. src/summarizer.py +2 -2
src/summarizer.py CHANGED
@@ -1,4 +1,4 @@
1
- from transformers import AutoTokenizer, AutoModelForSeq2SeqGeneration
2
  import torch
3
  import logging
4
 
@@ -16,7 +16,7 @@ class TextSummarizer:
16
 
17
  # Initialize tokenizer and model separately
18
  self.tokenizer = AutoTokenizer.from_pretrained(model_name)
19
- self.model = AutoModelForSeq2SeqGeneration.from_pretrained(model_name)
20
 
21
  # Move model to CPU and eval mode
22
  self.model = self.model.to(self.device)
 
1
+ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
2
  import torch
3
  import logging
4
 
 
16
 
17
  # Initialize tokenizer and model separately
18
  self.tokenizer = AutoTokenizer.from_pretrained(model_name)
19
+ self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
20
 
21
  # Move model to CPU and eval mode
22
  self.model = self.model.to(self.device)