Spaces:

ml6team
/

keyphrase-extraction

Running

App Files Files Community

DeDeckerThomas commited on May 12, 2022

Commit

55b038b

•

1 Parent(s): 860c19c

Fix last bugs

Browse files

Files changed (2) hide show

pipelines/__pycache__/keyphrase_generation_pipeline.cpython-39.pyc +0 -0
pipelines/keyphrase_generation_pipeline.py +3 -3

pipelines/__pycache__/keyphrase_generation_pipeline.cpython-39.pyc CHANGED Viewed

Binary files a/pipelines/__pycache__/keyphrase_generation_pipeline.cpython-39.pyc and b/pipelines/__pycache__/keyphrase_generation_pipeline.cpython-39.pyc differ

pipelines/keyphrase_generation_pipeline.py CHANGED Viewed

@@ -3,6 +3,7 @@ from transformers import (
     AutoModelForSeq2SeqLM,
     AutoTokenizer,
 )
 class KeyphraseGenerationPipeline(Text2TextGenerationPipeline):
@@ -17,14 +18,13 @@ class KeyphraseGenerationPipeline(Text2TextGenerationPipeline):
     def postprocess(self, model_outputs):
         results = super().postprocess(model_outputs=model_outputs)
-        print(results)
         return [
             [
-                keyphrase.strip()
                 for keyphrase in result.get("generated_text").split(
                     self.keyphrase_sep_token
                 )
-                if keyphrase != ""
             ]
             for result in results
         ][0]

     AutoModelForSeq2SeqLM,
     AutoTokenizer,
 )
+import string
 class KeyphraseGenerationPipeline(Text2TextGenerationPipeline):
     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]