Rename qa_bot2.py to qa_bot_chatgpt.py

qa_bot2.py → qa_bot_chatgpt.py

@@ -6,28 +6,12 @@ import os
 import torch
 from huggingface_hub import login
 import ast
-
+from openai import OpenAI
 class QAInfer:
     def __init__(self):
         torch.cuda.empty_cache()
-        qa_model_name = "mistralai/Mistral-7B-Instruct-v0.2"
-        bnb_config = BitsAndBytesConfig(
-                load_in_4bit=True,
-                bnb_4bit_quant_type="nf4",
-                bnb_4bit_compute_dtype="float16",
-                bnb_4bit_use_double_quant=False,
-            )
-        self.qa_tokenizer = AutoTokenizer.from_pretrained(qa_model_name)
-        self.qa_tokenizer.pad_token = self.qa_tokenizer.eos_token
-
-        self.qa_model = AutoModelForCausalLM.from_pretrained(
-                qa_model_name,
-                quantization_config=bnb_config,
-                torch_dtype=torch.bfloat16,
-                device_map="auto",
-                trust_remote_code=True,
-                cache_dir="cache/mistral"
-            )
+       
+        self.chatgpt_client = OpenAI(api_key="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")
 
     def extract_text_from_pdf(self, pdf_path):
         """Extract text from a PDF file."""
@@ -103,27 +87,20 @@ class QAInfer:
         {pdf_text}
 
         === Response ===
-        If the answer cannot be found in the context, please respond with "Answer not found in the context""""
+         If related answer is not found return 'Information not present in the pdf' and below it provide something related to the question"""
 
-        inputs = self.qa_tokenizer(prompt, return_tensors='pt', truncation=True, max_length=512)
-        pipe = pipeline(
-            "text-generation", 
-            model=self.qa_model, 
-            tokenizer=self.qa_tokenizer, 
-            torch_dtype=torch.bfloat16, 
-            device_map="auto"
-        )
-        sequences = pipe(
-            prompt,
-            do_sample=True,
-            max_new_tokens=200, 
-            temperature=0.7, 
-            top_k=50, 
-            top_p=0.95,
-            num_return_sequences=1,
+    
+
+        print(prompt)
+        completion = self.chatgpt_client.chat.completions.create(
+            model="gpt-3.5-turbo",
+            messages=[
+                {"role": "system", "content": "You are a expert PDF parser , go through the pdf and answer the question properly , if related answer is not found return 'Information not present in the pdf' and below it provide something related to the question"},
+                {"role": "user", "content": prompt }
+            ]
         )
-        answer = sequences[0]['generated_text']
-        return answer
+        return completion.choices[0].message.content
+
 
 if __name__ == '__main__':
     qa_infer = QAInfer()