hanzla commited on
Commit
52f9aa7
1 Parent(s): 4dc6167
Files changed (1) hide show
  1. src/pdfchatbot.py +5 -1
src/pdfchatbot.py CHANGED
@@ -53,6 +53,9 @@ class PDFChatBot:
53
  docs = text_splitter.split_documents(self.documents)
54
  self.vectordb = Chroma.from_documents(docs, self.embeddings)
55
  print("Vector store created")
 
 
 
56
 
57
  @spaces.GPU
58
  def create_organic_pipeline(self):
@@ -62,6 +65,7 @@ class PDFChatBot:
62
  model_kwargs={"torch_dtype": torch.bfloat16},
63
  device="cuda",
64
  )
 
65
  print("Model pipeline loaded")
66
 
67
  def get_organic_context(self, query):
@@ -77,7 +81,7 @@ class PDFChatBot:
77
  {"role": "system", "content": "From the the contained given below, answer the question of user \n " + self.current_context},
78
  {"role": "user", "content": query},
79
  ]
80
- prompt = self.pipeline.tokenizer.apply_chat_template(
81
  messages,
82
  tokenize=False,
83
  add_generation_prompt=True
 
53
  docs = text_splitter.split_documents(self.documents)
54
  self.vectordb = Chroma.from_documents(docs, self.embeddings)
55
  print("Vector store created")
56
+ @spaces.GPU
57
+ def load_tokenizer(self):
58
+ self.tokenizer = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
59
 
60
  @spaces.GPU
61
  def create_organic_pipeline(self):
 
65
  model_kwargs={"torch_dtype": torch.bfloat16},
66
  device="cuda",
67
  )
68
+ self.load_tokenizer()
69
  print("Model pipeline loaded")
70
 
71
  def get_organic_context(self, query):
 
81
  {"role": "system", "content": "From the the contained given below, answer the question of user \n " + self.current_context},
82
  {"role": "user", "content": query},
83
  ]
84
+ prompt = self.tokenizer.apply_chat_template(
85
  messages,
86
  tokenize=False,
87
  add_generation_prompt=True