Spaces:

open-nlp
/

Chris-lab

Sleeping

kz209 commited on Jul 30, 2024

Commit

6b6861a

1 Parent(s): c5b9462

change to summarization

Files changed (1) hide show

utils/model.py CHANGED Viewed

@@ -24,7 +24,7 @@ class Model(torch.nn.Module):
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
         self.name = model_name
         self.pipeline = transformers.pipeline(
-            "text-generation",
             model=model_name,
             tokenizer=self.tokenizer,
             torch_dtype=torch.bfloat16,

         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
         self.name = model_name
         self.pipeline = transformers.pipeline(
+            "summarization",
             model=model_name,
             tokenizer=self.tokenizer,
             torch_dtype=torch.bfloat16,