Refactoring

main.py

@@ -12,6 +12,8 @@ import time
 
 # Constants
 GREETING = "Hi there! I'm an AI agent that uses a [retrieval-augmented generation](https://en.wikipedia.org/wiki/Retrieval-augmented_generation) pipeline to answer questions about research by the Design Research Collective. And the best part is that I always cite my ssources! What can I tell you about today?"
+EMBEDDING_MODEL_NAME = "allenai-specter"
+LLM_MODEL_NAME = "Qwen/Qwen2-7B-Instruct"
 
 # Load the dataset and convert to pandas
 full_data = datasets.load_dataset("ccm/publications")["train"].to_pandas()
@@ -34,7 +36,7 @@ index.train(vectors)
 index.add(vectors)
 
 # Load the model for later use in embeddings
-model = sentence_transformers.SentenceTransformer("allenai-specter")
+model = sentence_transformers.SentenceTransformer(EMBEDDING_MODEL_NAME)
 
 # Define the search function
 def search(query: str, k: int) -> tuple[str]:
@@ -60,11 +62,10 @@ def search(query: str, k: int) -> tuple[str]:
 
 
 # Create an LLM pipeline that we can send queries to
-model_name = "Qwen/Qwen2-7B-Instruct"
-tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
+tokenizer = transformers.AutoTokenizer.from_pretrained(LLM_MODEL_NAME)
 streamer = transformers.TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 chatmodel = transformers.AutoModelForCausalLM.from_pretrained(
-    model_name,
+    LLM_MODEL_NAME,
     torch_dtype="auto",
     device_map="auto"
 )