Upload handler.py

handler.py

@@ -0,0 +1,69 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from langchain.llms import HuggingFacePipeline
+from langchain import PromptTemplate, LLMChain
+from typing import Dict, List, Any
+import holidays
+
+template = """{char_name}'s Persona: {char_persona}
+<START>
+{chat_history}
+{char_name}: {char_greeting}
+<END>
+{user_name}: {user_input}
+{char_name}: """
+
+class EndpointHandler():
+
+    def __init__(self, path=""):
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForCausalLM.from_pretrained(path, load_in_8bit = True, device_map = "auto")
+        self.local_llm = HuggingFacePipeline(
+            pipeline = pipeline(
+                "text-generation",
+                model = self.model,
+                tokenizer = self.tokenizer,
+                max_length = 2048,
+                temperature = 0.5,
+                top_p = 0.9,
+                top_k = 0,
+                repetition_penalty = 1.1,
+                pad_token_id = 50256,
+                num_return_sequences = 1
+            )
+        )
+        self.prompt_template = PromptTemplate(
+            template = template,
+            input_variables = [
+                "user_input",
+                "user_name",
+                "char_name",
+                "char_persona",
+                "char_greeting",
+                "chat_history"
+            ],
+            validate_template = True
+        )
+        self.llm_engine = LLMChain(
+            llm = self.local_llm,
+            prompt = self.prompt_template,
+            verbose = True
+        )
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        data args:
+            inputs (:obj: `str`)
+            date (:obj: `str`)
+        Return:
+            A :obj:`list` | `dict`: will be serialized and returned
+        """
+        inputs = data.pop("inputs", data)
+
+        return self.llm_engine.predict(
+            user_input = inputs["user_input"],
+            user_name = inputs["user_name"],
+            char_name = inputs["char_name"],
+            char_persona = inputs["char_persona"],
+            char_greeting = inputs["char_greeting"],
+            chat_history = inputs["chat_history"]
+        ).split("\n",1)[0]