Update main.py

main.py

@@ -1,7 +1,9 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import os
+os.environ["CUDA_VISIBLE_DEVICES"]="0"
 
+device = torch.device("cuda")
 
 model = AutoModelForCausalLM.from_pretrained(
     "E-Hospital/open-orca-platypus-2-lora-medical",
@@ -11,7 +13,7 @@ model = AutoModelForCausalLM.from_pretrained(
 tokenizer = AutoTokenizer.from_pretrained("Open-Orca/OpenOrca-Platypus2-13B", trust_remote_code=True)
 
 def ask_bot(question):
-  input_ids = tokenizer.encode(question, return_tensors="pt").to('cuda')
+  input_ids = tokenizer.encode(question, return_tensors="pt").to(device)
   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)
@@ -53,7 +55,7 @@ class CustomLLM(LLM):
         if stop is not None:
             raise ValueError("stop kwargs are not permitted.")
 
-        input_ids = tokenizer.encode(prompt, return_tensors="pt").to('cuda')
+        input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
         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)
@@ -67,15 +69,6 @@ class CustomLLM(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
-
-
 class DbHandler():
     def __init__(self):
         self.db_con = mysql.connector.connect(