Update chat.py

chat.py

@@ -1,23 +1,13 @@
 from peft import AutoPeftModelForCausalLM
-from transformers import GenerationConfig
-from transformers import AutoTokenizer
+from transformers import GenerationConfig, AutoTokenizer
 import torch
 import os
-import os
 
+# Set huggingface cache directory
 os.environ["HF_HOME"] = "./hf_home"
 os.environ["TRANSFORMERS_CACHE"] = "./hf_home/transformers"
-os.makedirs("./hf_home/transformers", exist_ok=True)
-
-
-# Set Hugging Face cache directory to a folder you have access to
-os.environ["HF_HOME"] = "/data"
-os.environ["TRANSFORMERS_CACHE"] = "/data/transformers"
-
-# Create cache directory if doesn't exist
-os.makedirs("/data/transformers", exist_ok=True)
-
 
+# Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("Vasanth/mistral-finetuned-alpaca")
 
 model = AutoPeftModelForCausalLM.from_pretrained(
@@ -25,7 +15,8 @@ model = AutoPeftModelForCausalLM.from_pretrained(
     low_cpu_mem_usage=True,
     return_dict=True,
     torch_dtype=torch.float16,
-    device_map="cuda")
+    device_map="auto"
+)
 
 generation_config = GenerationConfig(
     do_sample=True,
@@ -36,7 +27,8 @@ generation_config = GenerationConfig(
 )
 
 def chatbot(message):
-    input_str = "###Human: " + message + " ###Assistant: "
-    inputs = tokenizer(input_str, return_tensors="pt").to("cuda")
+    input_str = f"###Human: {message} ###Assistant: "
+    inputs = tokenizer(input_str, return_tensors="pt").to(model.device)
     outputs = model.generate(**inputs, generation_config=generation_config)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True).replace(input_str, '')
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True).replace(input_str, "")
+    return response.strip()