Upload folder using huggingface_hub

handler.py

@@ -8,17 +8,40 @@ from transformers import (
     RobertaTokenizer,
     RobertaForSequenceClassification,
 )
+import nltk
+from nltk.tokenize import sent_tokenize
+nltk.download('punkt')
 
-class DipperParaphraser(object):
-    def __init__(self, model="", verbose=True):
-        time1 = time.time()
-        self.tokenizer = T5Tokenizer.from_pretrained('google/t5-v1_1-xxl')
-        self.model = T5ForConditionalGeneration.from_pretrained(model, device_map="auto", load_in_8bit=True)
-        if verbose:
-            print(f"{model} model loaded in {time.time() - time1}")
-        # self.model.cuda()
-        self.model.eval()
 
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = T5ForConditionalGeneration.from_pretrained(path, device_map="auto", load_in_8bit=True)
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        """
+       data args:
+            inputs (:obj: `str`)
+            date (:obj: `str`)
+      Return:
+            A :obj:`list` | `dict`: will be serialized and returned
+        """
+        input_text = data.get("input_text", "")
+        lex_diversity = data.get("lex_diversity", 80)
+        order_diversity = data.get("order_diversity", 20)
+        prefix = data.get("prefix", "")
+        prediction = self.paraphrase(
+            input_text, 
+            lex_diversity, 
+            order_diversity, 
+            prefix=prefix, 
+            do_sample=True, 
+            top_p=0.75,  
+            max_length=512
+        )
+
+        prediction = {'prediction': prediction}
+        return prediction
     def paraphrase(self, input_text, lex_diversity, order_diversity, prefix="", sent_interval=3, **kwargs):
         """Paraphrase a text using the DIPPER model.
 
@@ -55,34 +78,4 @@ class DipperParaphraser(object):
             prefix += " " + outputs[0]
             output_text += " " + outputs[0]
 
-        return output_text
-
-
-class EndpointHandler:
-    def __init__(self, path=""):
-        self.pipeline = DipperParaphraser(model=path)
-
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        """
-       data args:
-            inputs (:obj: `str`)
-            date (:obj: `str`)
-      Return:
-            A :obj:`list` | `dict`: will be serialized and returned
-        """
-        input_text = data.get("input_text", "")
-        lex_diversity = data.get("lex_diversity", 80)
-        order_diversity = data.get("order_diversity", 20)
-        prefix = data.get("prefix", "")
-        prediction = self.pipeline.paraphrase(
-            input_text, 
-            lex_diversity, 
-            order_diversity, 
-            prefix=prefix, 
-            do_sample=True, 
-            top_p=0.75,  
-            max_length=512
-        )
-
-        prediction = {'prediction': prediction}
-        return prediction
+        return output_text