Update handler.py

handler.py

@@ -11,7 +11,9 @@ tokenizer.push_to_hub(repo_name="text-summary-gpt2-short", repo_id="Lin0He/text-
 '''
 import torch
 from typing import Dict, List, Any
-from transformers import pipeline, AutoModel, AutoTokenizer
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 def topk(probs, n=9):
     # The scores are initially softmaxed to convert to probabilities
@@ -65,8 +67,8 @@ def model_infer(model, tokenizer, review, max_length=300):
 def predict(text, model, tokenizer):
     result_text = []
     for i in range(6):
-        summary = model_infer(model, tokenizer, input+"TL;DR").strip()
-        result_text.append(summary[len(input)+5:])
+        summary = model_infer(model, tokenizer, text+"TL;DR").strip()
+        result_text.append(summary[len(text)+5:])
     return sorted(result_text, key=len)[3]
     #print("summary:", sorted(result_text, key=len)[3])
 
@@ -74,7 +76,7 @@ class EndpointHandler():
     def __init__(self, path="Lin0He/text-summary-gpt2-short"):
         # load model and tokenizer from path
         self.tokenizer = AutoTokenizer.from_pretrained(path)
-        self.model = AutoModel.from_pretrained(path)
+        self.model = AutoModelForCausalLM.from_pretrained(path)
 
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: