Skip to content

Commit

Permalink
Fix summary pipeline issue with transformers>=4.19.0, closes #278
Browse files Browse the repository at this point in the history
  • Loading branch information
davidmezzetti committed May 12, 2022
1 parent 1a8a8f1 commit 248acc1
Showing 1 changed file with 26 additions and 5 deletions.
31 changes: 26 additions & 5 deletions src/python/txtai/pipeline/text/summary.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,11 +39,8 @@ def __call__(self, text, minlength=None, maxlength=None, workers=0):
texts = text if isinstance(text, list) else [text]
params = [(x, text if len(text) >= check else None) for x, text in enumerate(texts)]

kwargs = {"truncation": True}
if minlength:
kwargs["min_length"] = minlength
if maxlength:
kwargs["max_length"] = maxlength
# Build keyword arguments
kwargs = self.args(minlength, maxlength)

inputs = [text for _, text in params if text]
if inputs:
Expand Down Expand Up @@ -74,3 +71,27 @@ def clean(self, text):
text = text.strip()

return text

def args(self, minlength, maxlength):
"""
Builds keyword arguments.
Args:
minlength: minimum length for summary
maxlength: maximum length for summary
Returns:
keyword arguments
"""

kwargs = {"truncation": True}
if minlength:
kwargs["min_length"] = minlength
if maxlength:
kwargs["max_length"] = maxlength

# Default minlength if not provided or it's bigger than maxlength
if "min_length" not in kwargs or kwargs["min_length"] > kwargs["max_length"]:
kwargs["min_length"] = kwargs["max_length"]

return kwargs

0 comments on commit 248acc1

Please sign in to comment.