Update main.py

main.py

@@ -26,7 +26,7 @@ EXAMPLE_QUERIES = [
     "What is known about different modes for human-AI teaming?",
 ]
 EMBEDDING_MODEL_NAME = "allenai-specter"
-LLM_MODEL_NAME = "Qwen/Qwen2-7B-Instruct-AWQ"
+LLM_MODEL_NAME = "Qwen/Qwen2-7B-Instruct"
 
 # Load the dataset and convert to pandas
 data = pandas.read_parquet("hf://datasets/ccm/publications/data/train-00000-of-00001.parquet")
@@ -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, torch_dtype=torch.float16)
+tokenizer = transformers.AutoTokenizer.from_pretrained(LLM_MODEL_NAME)
 streamer = transformers.TextIteratorStreamer(
-    tokenizer, skip_prompt=True, skip_special_tokens=True, torch_dtype=torch.float16
+    tokenizer, skip_prompt=True, skip_special_tokens=True
 )
 chatmodel = transformers.AutoModelForCausalLM.from_pretrained(
-    LLM_MODEL_NAME, device_map="auto", torch_dtype=torch.float16
+    LLM_MODEL_NAME, device_map="auto", torch_dtype="auto"
 )
 
 # Create a FAISS index for fast similarity search