Update pipeline.py

pipeline.py

@@ -30,7 +30,7 @@ def topk(probs, n=9):
     tokenId = topIx[choice][0]
     return int(tokenId)
 
-def model_infer(model, tokenizer, review, max_length=300):
+def model_infer(model, tokenizer, review, max_length=10):
     result_text = []
     for i in range(6):
     
@@ -61,13 +61,14 @@ def model_infer(model, tokenizer, review, max_length=300):
     
                 # If the chosen token is EOS, return the result
                 if res_id == tokenizer.eos_token_id:
-                    return tokenizer.decode(result)
+                    result_text.append(tokenizer.decode(result)[len(review):])
+                    break
                 else: # Append to the sequence
                     result.append(res_id)
 
         # IF no EOS is generated, return after the max_len
-        result_text.append(tokenizer.decode(result)) 
-    return sorted(result_text, key=len)[3]
+        #result_text.append(tokenizer.decode(result)) 
+    return sorted(result_text, key=len)[4]
 
 class PreTrainedPipeline():
     def __init__(self, path=""):