Updated rag.py

rag.py

@@ -10,10 +10,9 @@ class RAG:
     def __init__(self):
         self.model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        local_dir = "llm_models/"
 
         self.embedding_model_name = "all-mpnet-base-v2"
-        self.embeddings_filename = "data/embeddings.csv"
+        self.embeddings_filename = "embeddings.csv"
         
         self.data_pd = pd.read_csv(self.embeddings_filename)
         self.data_dict = pd.read_csv(self.embeddings_filename).to_dict(orient='records')
@@ -23,11 +22,9 @@ class RAG:
         # Embedding model
         self.embedding_model = SentenceTransformer(model_name_or_path = self.embedding_model_name,device = self.device)
         # Tokenizer
-        self.tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=self.model_id,
-                                                      cache_dir = local_dir)
+        self.tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=self.model_id)
         # LLM
-        self.llm_model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=self.model_id,
-                                                         cache_dir = local_dir).to(self.device)
+        self.llm_model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=self.model_id).to(self.device)
 
     def get_embeddings(self) -> list:
         """Returns the embeddings from the csv file"""