Update handler.py

handler.py

@@ -16,22 +16,21 @@ class EndpointHandler():
         messages = request_inputs["messages"]
         char_name = request_inputs["char_name"]
         user_name = request_inputs["user_name"]
+        chats_curled = request_inputs["chats_curled"]
         template = self.default_template
-        user_input = "\n".join([
+        user_input = [
             "{name}: {message}".format(
                 name = char_name if (id["role"] == "AI") else user_name,
                 message = id["message"].strip()
             ) for id in messages
-        ])
-        prompt = template.format(
-            char_name = char_name,
-            user_name = user_name,
-            user_input = user_input
-        )
-        input_ids = self.tokenizer(
-            prompt + f"\n{char_name}:",
-            return_tensors = "pt"
-        ).to("cuda")
+        ]
+        while True:
+            prompt = template.format(char_name = char_name, user_name = user_name, user_input = "\n".join([user_input]))
+            input_ids = self.tokenizer(prompt + f"\n{char_name}:", return_tensors = "pt").to("cuda")
+            if input_ids.input_ids.size(1) > 2048:
+                chats_curled += 1
+                user_input = user_input[chats_curled*2:]
+            else: break
         encoded_output = self.model.generate(
             input_ids["input_ids"],
             max_new_tokens = 50,