Update main.py

main.py

@@ -1,38 +1,27 @@
-# -*- coding: utf-8 -*-
-"""TestAPI.ipynb
-
-Automatically generated by Colaboratory.
-
-Original file is located at
-    https://colab.research.google.com/drive/1WToaz7kQoFpI0_M8j6uWPigBrKlkL4ml
-"""
-
-
-
-from transformers import AutoTokenizer,AutoModelForCausalLM
-import os
-os.environ["CUDA_VISIBLE_DEVICES"]="0"
-import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-from typing import List
-
+model = AutoModelForCausalLM.from_pretrained(
+    "E-Hospital/open-orca-platypus-2-lora-medical",
+    trust_remote_code=True,
+    device_map = "auto",
+)
+tokenizer = AutoTokenizer.from_pretrained("Open-Orca/OpenOrca-Platypus2-13B", trust_remote_code=True)
 
-from peft import PeftModel
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from typing import Any, List, Mapping, Optional
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from langchain.callbacks.manager import CallbackManagerForLLMRun
-from langchain.llms.base import LLM
+def ask_bot(question):
+  input_ids = tokenizer.encode(question, return_tensors="pt").to('cuda')
+  with torch.no_grad():
+      output = model.generate(input_ids, max_length=500, num_return_sequences=1, do_sample=True, top_k=50)
+  generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+  response = generated_text.split("->:")[-1]
+  return response
 
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+ask_bot("I have diabetes. What should I do?")
 
 import mysql.connector
 import re
 from datetime import datetime
-
+from typing import Any, List, Mapping, Optional
 
 from langchain.memory import ConversationBufferMemory
 from typing import Any, List, Mapping, Optional
@@ -46,28 +35,6 @@ from langchain.memory import ConversationSummaryBufferMemory
 
 from langchain.memory import ConversationSummaryMemory
 
-
-model_name = "Open-Orca/OpenOrca-Platypus2-13B"
-
-
-tokenizer = AutoTokenizer.from_pretrained(
-    model_name, trust_remote_code=True)
-
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    trust_remote_code=True,
-    load_in_8bit = True,
-    device_map = "auto",
-)
-
-
-
-model = PeftModel.from_pretrained(model, "teslalord/open-orca-platypus-2-medical")
-
-model = model.merge_and_unload()
-
-
-
 class CustomLLM(LLM):
     n: int
     # custom_model: llm  # Replace with the actual type of your custom model
@@ -137,7 +104,7 @@ class DbHandler():
 
     def get_history(self, patient_id):
         try:
-            query = f"SELECT * FROM chatbot_conversation WHERE patient_id = '{patient_id}' ORDER BY timestamp ASC;"
+            query = f"SELECT * FROM chatbot_conversation WHERE patient_id = {patient_id} ORDER BY timestamp ASC;"
             self.cursorObject.execute(query)
             data = self.cursorObject.fetchall()
             return data
@@ -159,8 +126,6 @@ def get_conversation_history(db, patient_id):
         return conversations[-1][5]
     return ""
 
-
-
 llm = CustomLLM(n=10)
 app = FastAPI()
 
@@ -201,4 +166,3 @@ def chatbot(patient_id, user_data: dict=None):
       return {"response": response}
     finally:
       db.close_db()
-