Update main.py

main.py

@@ -27,7 +27,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"
+LLM_MODEL_NAME = "microsoft/Phi-3-small-8k-instruct-onnx-cuda" #"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")
@@ -43,12 +43,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, trust_remote_code=True)
 streamer = transformers.TextIteratorStreamer(
     tokenizer, skip_prompt=True, skip_special_tokens=True
 )
 chatmodel = transformers.AutoModelForCausalLM.from_pretrained(
-    LLM_MODEL_NAME, device_map="auto", torch_dtype="auto"
+    LLM_MODEL_NAME, device_map="auto", torch_dtype="auto", trust_remote_code=True
 )
 
 # Create a FAISS index for fast similarity search