Mojo3 commited on
Commit
7b00046
·
verified ·
1 Parent(s): b9fe0e0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +9 -2
app.py CHANGED
@@ -10,6 +10,13 @@ from langchain.vectorstores import Chroma
10
  from langchain.docstore.document import Document as Document2
11
  from langchain_community.embeddings import HuggingFaceEmbeddings
12
 
 
 
 
 
 
 
 
13
 
14
  docs_folder = "./converted_docs"
15
 
@@ -130,9 +137,9 @@ class RAGPipeline:
130
  self.prompt_template = PromptTemplate.from_template(self._get_template())
131
 
132
  # Load model and tokenizer
133
- self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
134
  self.model = AutoModelForCausalLM.from_pretrained(
135
- self.model_name, torch_dtype=torch.bfloat16, device_map="auto"
136
  )
137
 
138
  def _get_template(self):
 
10
  from langchain.docstore.document import Document as Document2
11
  from langchain_community.embeddings import HuggingFaceEmbeddings
12
 
13
+ from huggingface_hub import HfFolder
14
+
15
+ # Load token from environment variable
16
+ token = os.getenv("token")
17
+
18
+ # Save the token to Hugging Face's system directory
19
+ HfFolder.save_token(token)
20
 
21
  docs_folder = "./converted_docs"
22
 
 
137
  self.prompt_template = PromptTemplate.from_template(self._get_template())
138
 
139
  # Load model and tokenizer
140
+ self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, token=token)
141
  self.model = AutoModelForCausalLM.from_pretrained(
142
+ self.model_name, torch_dtype=torch.bfloat16, device_map="auto", token=token
143
  )
144
 
145
  def _get_template(self):