Skip to content

Commit 7ee7c97

Browse files
author
Marcin Kardas
committed
Add return_full_text option to generate
1 parent e3e3448 commit 7ee7c97

File tree

2 files changed

+6
-2
lines changed

2 files changed

+6
-2
lines changed

galai/model.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -223,6 +223,7 @@ def generate(
223223
penalty_alpha=None,
224224
num_beams=1,
225225
num_return_sequences=1,
226+
return_full_text=True,
226227
) -> Union[str, List[str], List[List[str]]]:
227228
"""
228229
Generates text using the model
@@ -301,8 +302,11 @@ def generate(
301302
**options
302303
)
303304

305+
out_tokens = out['sequences']
306+
if not return_full_text:
307+
out_tokens = out_tokens[:, input_v.shape[1]:]
304308
# we keep special tokens such as [START_REF] or <work>
305-
decoded = self.tokenizer.batch_decode(out['sequences'], skip_special_tokens=False)
309+
decoded = self.tokenizer.batch_decode(out_tokens, skip_special_tokens=False)
306310
# so we manually remove </s> and <pad>
307311
decoded = [
308312
text.replace(self.tokenizer.eos_token, "").replace(self.tokenizer.pad_token, "")

setup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from setuptools import setup, find_packages
22

33
PACKAGE_NAME = 'galai'
4-
VERSION = "1.1.4"
4+
VERSION = "1.1.5"
55
DESCRIPTION = "API for the GALACTICA model"
66
KEYWORDS = "Scientific Intelligence"
77
URL = 'https://github.com/paperswithcode/galai'

0 commit comments

Comments
 (0)