Rahatara commited on
Commit
3994ed5
1 Parent(s): f361940

Update pdfchatbot.py

Browse files
Files changed (1) hide show
  1. pdfchatbot.py +10 -4
pdfchatbot.py CHANGED
@@ -11,13 +11,16 @@ from langchain.document_loaders import PyPDFLoader
11
  from langchain.prompts import PromptTemplate
12
  from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
13
 
 
 
 
14
  class PDFChatBot:
15
  def __init__(self, config_path="config.yaml"):
16
  """
17
  Initialize the PDFChatBot instance.
18
 
19
  Parameters:
20
- config_path (str): Path to the configuration file (default is "../config.yaml").
21
  """
22
  self.processed = False
23
  self.page = 0
@@ -94,7 +97,10 @@ class PDFChatBot:
94
  """
95
  Load the tokenizer from Hugging Face and set in the config file.
96
  """
97
- self.tokenizer = AutoTokenizer.from_pretrained(self.config.get("autoTokenizer"))
 
 
 
98
 
99
  def load_model(self):
100
  """
@@ -104,7 +110,7 @@ class PDFChatBot:
104
  self.config.get("autoModelForCausalLM"),
105
  device_map='auto',
106
  torch_dtype=torch.float32,
107
- token=True,
108
  load_in_8bit=False
109
  )
110
 
@@ -190,4 +196,4 @@ class PDFChatBot:
190
  page = doc[self.page]
191
  pix = page.get_pixmap(matrix=fitz.Matrix(300 / 72, 300 / 72))
192
  image = Image.frombytes('RGB', [pix.width, pix.height], pix.samples)
193
- return image
 
11
  from langchain.prompts import PromptTemplate
12
  from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
13
 
14
+ # Store the Hugging Face token in a variable
15
+ HUGGINGFACE_TOKEN = gr.Textbox()
16
+
17
  class PDFChatBot:
18
  def __init__(self, config_path="config.yaml"):
19
  """
20
  Initialize the PDFChatBot instance.
21
 
22
  Parameters:
23
+ config_path (str): Path to the configuration file (default is "config.yaml").
24
  """
25
  self.processed = False
26
  self.page = 0
 
97
  """
98
  Load the tokenizer from Hugging Face and set in the config file.
99
  """
100
+ self.tokenizer = AutoTokenizer.from_pretrained(
101
+ self.config.get("autoTokenizer"),
102
+ use_auth_token=HUGGINGFACE_TOKEN
103
+ )
104
 
105
  def load_model(self):
106
  """
 
110
  self.config.get("autoModelForCausalLM"),
111
  device_map='auto',
112
  torch_dtype=torch.float32,
113
+ use_auth_token=HUGGINGFACE_TOKEN,
114
  load_in_8bit=False
115
  )
116
 
 
196
  page = doc[self.page]
197
  pix = page.get_pixmap(matrix=fitz.Matrix(300 / 72, 300 / 72))
198
  image = Image.frombytes('RGB', [pix.width, pix.height], pix.samples)
199
+ return image