Update app.py

app.py

@@ -16,23 +16,34 @@ tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 model = PeftModel.from_pretrained(model, peft_model_id)
 
 
+def make_inference(problem):
+    batch = tokenizer(
+        f"User: {product_name} ,
+        return_tensors="pt",
+    )
 
-def make_inference(conversation):
-    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
+    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):
+#     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__":