Update app.py

app.py

@@ -16,30 +16,18 @@ tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 model = PeftModel.from_pretrained(model, peft_model_id)
 
 
-# def make_inference(problem, transcript):
-#     batch = tokenizer(
-#         f"### Problem:\n{problem}: \n\n### Transcript:",
-#         return_tensors="pt",
-#     )
-
-#     with torch.cuda.amp.autocast():
-#         output_tokens = model.generate(**batch, max_new_tokens=50)
-
-#     return tokenizer.decode(output_tokens[0], skip_special_tokens=True)
-
-def make_inference(conversation):
-    batch = tokenizer(
-        f"User:{conversation}\n",
-        return_tensors="pt",
-    )
 
+def make_inference(conversations):
+    context = ""
+    for conversation in conversations:
+        context += f"{conversation}\n\n"
+    prompt = f"### Conversation:\n{context}"
+    batch = tokenizer(prompt, return_tensors="pt")
     with torch.cuda.amp.autocast():
         output_tokens = model.generate(**batch, max_new_tokens=50)
-
     response = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
-
-    #updated_conversation = f"{conversation}\n\n{response}"
-    return response
+    updated_conversation = f"{context}\n{response}"
+    return updated_conversation
 
 
 if __name__ == "__main__":