Update main.py

main.py

@@ -42,12 +42,12 @@ data.reset_index(inplace=True)
 model = sentence_transformers.SentenceTransformer(EMBEDDING_MODEL_NAME)
 
 # Create an LLM pipeline that we can send queries to
-tokenizer = transformers.AutoTokenizer.from_pretrained(LLM_MODEL_NAME)
+tokenizer = transformers.AutoTokenizer.from_pretrained(LLM_MODEL_NAME, torch_dtype=torch.float16)
 streamer = transformers.TextIteratorStreamer(
-    tokenizer, skip_prompt=True, skip_special_tokens=True
+    tokenizer, skip_prompt=True, skip_special_tokens=True, torch_dtype=torch.float16
 )
 chatmodel = transformers.AutoModelForCausalLM.from_pretrained(
-    LLM_MODEL_NAME, torch_dtype="auto", device_map="auto"
+    LLM_MODEL_NAME, device_map="auto", torch_dtype=torch.float16
 )
 
 # Create a FAISS index for fast similarity search