Update app.py

app.py

@@ -15,35 +15,35 @@ tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 # Load the Lora model
 model = PeftModel.from_pretrained(model, peft_model_id)
 
-# def make_inference(problem):
-#     batch = tokenizer(
-#         problem,
-#         return_tensors="pt",
-#     )
+def make_inference(problem, answer):
+    batch = tokenizer(
+        problem,
+        return_tensors="pt",
+    )
 
-#     with torch.cuda.amp.autocast():
-#         output_tokens = model.generate(**batch, max_new_tokens=50)
+    with torch.cuda.amp.autocast():
+        output_tokens = model.generate(**batch, max_new_tokens=50)
 
 
 
 
 
-def make_inference(conversation, response):
-    conversation_history = conversation
-    response = ""
-    while True:
-        batch = tokenizer(
-            f"### Problem:\n{conversation_history}\n{response}",
-            return_tensors="pt",
-        )
-        with torch.cuda.amp.autocast():
-            output_tokens = model.generate(**batch, max_new_tokens=50)
-        new_response = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
-        if new_response.strip() == "":
-            break
-        response = f"\n{new_response}"
-        conversation_history += response
-    return conversation_history
+# def make_inference(conversation, response):
+#     conversation_history = conversation
+#     response = ""
+#     while True:
+#         batch = tokenizer(
+#             f"### Problem:\n{conversation_history}\n{response}",
+#             return_tensors="pt",
+#         )
+#         with torch.cuda.amp.autocast():
+#             output_tokens = model.generate(**batch, max_new_tokens=50)
+#         new_response = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
+#         if new_response.strip() == "":
+#             break
+#         response = f"\n{new_response}"
+#         conversation_history += response
+#     return conversation_history
 
 
 if __name__ == "__main__":