Update extract_text_from_pdf.py

extract_text_from_pdf.py

@@ -12,7 +12,7 @@ import spaces
 
 warnings.filterwarnings('ignore')
 
-@spaces.GPU
+#@spaces.GPU
 class PDFTextExtractor:
     """
     A class to handle PDF text extraction and preprocessing for podcast preparation.
@@ -28,19 +28,21 @@ class PDFTextExtractor:
             model_name (str): Name of the model to use for text processing.
         """
 
-        model_name="meta-llama/Llama-3.2-1B-Instruct"
+        #model_name="meta-llama/Llama-3.2-1B-Instruct"
 
         self.pdf_path = pdf_path
         self.output_path = output_path
         self.max_chars = 100000
         self.chunk_size = 1000
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        
+        #self.device = "cuda" if torch.cuda.is_available() else "cpu"
         
         # Initialize model and tokenizer
-        self.accelerator = Accelerator()
-        self.model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16,use_safetensors=True,device_map=self.device)
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_safetensors=True)
-        self.model, self.tokenizer = self.accelerator.prepare(self.model, self.tokenizer)
+        # self.accelerator = Accelerator()
+        # self.model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16,use_safetensors=True,device_map=self.device)
+        # self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_safetensors=True)
+        # self.model, self.tokenizer = self.accelerator.prepare(self.model, self.tokenizer)
+        self.model_name="llama3-8b-8192"
 
         
         # System prompt for text processing
@@ -135,13 +137,24 @@ class PDFTextExtractor:
             {"role": "user", "content": text_chunk}
         ]
         
-        prompt = self.tokenizer.apply_chat_template(conversation, tokenize=False)
-        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+        # prompt = self.tokenizer.apply_chat_template(conversation, tokenize=False)
+        # inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+        
+        # with torch.no_grad():
+        #     output = self.model.generate(**inputs, temperature=0.7, top_p=0.9, max_new_tokens=512)
+        
+        # processed_text = self.tokenizer.decode(output[0], skip_special_tokens=True)[len(prompt):].strip()
+        client = Groq(
+            api_key=os.environ.get("GROQ_API_KEY"),
+        )
+
+        chat_completion = client.chat.completions.create(
+            messages=conversation,
+            model=self.model_name,
+        )
         
-        with torch.no_grad():
-            output = self.model.generate(**inputs, temperature=0.7, top_p=0.9, max_new_tokens=512)
+        processed_text = self.extract_tuple(chat_completion.choices[0].message.content)
         
-        processed_text = self.tokenizer.decode(output[0], skip_special_tokens=True)[len(prompt):].strip()
         return processed_text