Spaces:

open-nlp
/

Chris-lab

Sleeping

kz209 commited on Sep 3, 2024

Commit

8f783e1

•

1 Parent(s): d6563b9

update

Files changed (1) hide show

utils/model.py CHANGED Viewed

@@ -108,4 +108,4 @@ class Model(torch.nn.Module):
             temperature=temp,
             eos_token_id=self.tokenizer.eos_token_id,
         )
-        return self.tokenizer.batch_decode(outputs, skip_special_tokens=True)

             temperature=temp,
             eos_token_id=self.tokenizer.eos_token_id,
         )
+        return self.tokenizer.batch_decode(outputs[:, input_ids.shape[1]:], skip_special_tokens=True)