Update pdfchatbot.py
Browse files- 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 "
|
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(
|
|
|
|
|
|
|
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 |
-
|
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
|