Update src/summarizer.py
Browse files- src/summarizer.py +2 -2
src/summarizer.py
CHANGED
|
@@ -1,4 +1,4 @@
|
|
| 1 |
-
from transformers import AutoTokenizer,
|
| 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 =
|
| 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)
|