Update handler.py

handler.py

@@ -31,47 +31,43 @@ def topk(probs, n=9):
     return int(tokenId)
 
 def model_infer(model, tokenizer, review, max_length=300):
-    # Preprocess the init token (task designator)
-    review_encoded = tokenizer.encode(review)
-    result = review_encoded
-    initial_input = torch.tensor(review_encoded).unsqueeze(0).to(device)
-
-    with torch.set_grad_enabled(False):
-        # Feed the init token to the model
-        output = model(initial_input)
-
-        # Flatten the logits at the final time step
-        logits = output.logits[0,-1]
-
-        # Make a top-k choice and append to the result
-        #choices = [topk(logits) for i in range(5)]
-        choices = topk(logits)
-        result.append(choices)
-        
-        # For max_length times:
-        for _ in range(max_length):
-            # Feed the current sequence to the model and make a choice
-            input = torch.tensor(result).unsqueeze(0).to(device)
-            output = model(input)
+    result_text = []
+    for i in range(6):
+    
+        # Preprocess the init token (task designator)
+        review_encoded = tokenizer.encode(review)
+        result = review_encoded
+        initial_input = torch.tensor(review_encoded).unsqueeze(0).to(device)
+    
+        with torch.set_grad_enabled(False):
+            # Feed the init token to the model
+            output = model(initial_input)
+    
+            # Flatten the logits at the final time step
             logits = output.logits[0,-1]
-            res_id = topk(logits)
-
-            # If the chosen token is EOS, return the result
-            if res_id == tokenizer.eos_token_id:
-                return tokenizer.decode(result)
-            else: # Append to the sequence
-                result.append(res_id)
+    
+            # Make a top-k choice and append to the result
+            #choices = [topk(logits) for i in range(5)]
+            choices = topk(logits)
+            result.append(choices)
+            
+            # For max_length times:
+            for _ in range(max_length):
+                # Feed the current sequence to the model and make a choice
+                input = torch.tensor(result).unsqueeze(0).to(device)
+                output = model(input)
+                logits = output.logits[0,-1]
+                res_id = topk(logits)
+    
+                # If the chosen token is EOS, return the result
+                if res_id == tokenizer.eos_token_id:
+                    return tokenizer.decode(result)
+                else: # Append to the sequence
+                    result.append(res_id)
 
     # IF no EOS is generated, return after the max_len
-    return tokenizer.decode(result)
-
-def predict(text, model, tokenizer):
-    result_text = []
-    for i in range(6):
-        summary = model_infer(model, tokenizer, text+"TL;DR").strip()
-        result_text.append(summary[len(text)+5:])
+    result_text.append(tokenizer.decode(result)) 
     return sorted(result_text, key=len)[3]
-    #print("summary:", sorted(result_text, key=len)[3])
 
 class EndpointHandler():
     def __init__(self, path=""):
@@ -80,17 +76,9 @@ class EndpointHandler():
         self.model = AutoModelForCausalLM.from_pretrained("Lin0He/text-summary-gpt2-short")
 
 
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         # process input
         inputs = data.pop("inputs", data)
         # process input text
-        prediction = predict(inputs, self.model, self.tokenizer)
+        prediction = model_infer( self.model, self.tokenizer，inputs+"TL;DR")
         return prediction
-        
-
-        
-'''
-predictor = pipeline("summarization", model = model, tokenizer = tokenizer)
-result = predictor("Input text for prediction")
-print(result)
-'''