Update code/inference.py

code/inference.py

@@ -20,17 +20,23 @@ def predict_fn(data, load_list):
     messages = request_inputs["messages"]
     char_name = request_inputs["char_name"]
     user_name = request_inputs["user_name"]
+    chats_curled = request_inputs["chats_curled"]
     user_input = [
         "{name}: {message}".format(
             name = char_name if (id["role"] == "AI") else user_name,
             message = id["message"].strip()
         ) for id in messages
     ]
-    user_input = "\n".join([user_input])
-    prompt = template.format(char_name = char_name, user_name = user_name, user_input = user_input)
 
-    # tokenize the model input, generate and decode output
-    input_ids = tokenizer(prompt + f"\n{char_name}:", return_tensors = "pt").to("cuda")
+    # Tokenize the model input
+    while True:
+        prompt = template.format(char_name = char_name, user_name = user_name, user_input = "\n".join([user_input]))
+        input_ids = 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 = model.generate(
         input_ids["input_ids"],
         max_new_tokens = 50,
@@ -54,5 +60,6 @@ def predict_fn(data, load_list):
     except Exception: pass
     return {
         "role": "AI",
-        "message": decoded_output
+        "message": decoded_output,
+        "chats_curled": chats_curled
     }