Spaces:
Running
Running
File size: 1,025 Bytes
e4f39c4 55b038b e4f39c4 55b038b e4f39c4 55b038b e4f39c4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 |
from transformers import (
Text2TextGenerationPipeline,
AutoModelForSeq2SeqLM,
AutoTokenizer,
)
import string
class KeyphraseGenerationPipeline(Text2TextGenerationPipeline):
def __init__(self, model, keyphrase_sep_token=";", *args, **kwargs):
super().__init__(
model=AutoModelForSeq2SeqLM.from_pretrained(model),
tokenizer=AutoTokenizer.from_pretrained(model),
*args,
**kwargs
)
self.keyphrase_sep_token = keyphrase_sep_token
def postprocess(self, model_outputs):
results = super().postprocess(model_outputs=model_outputs)
return [
[
keyphrase.strip().translate(str.maketrans('', '', string.punctuation))
for keyphrase in result.get("generated_text").split(
self.keyphrase_sep_token
)
if keyphrase.translate(str.maketrans('', '', string.punctuation)) != ""
]
for result in results
][0]
|