Update parse.py

parse.py

@@ -1,9 +1,6 @@
-from groq import Groq
+from langchain_ollama import OllamaLLM
+from langchain_core.prompts import ChatPromptTemplate
 
-# Initialize Groq client with the API key directly
-client = Groq(api_key="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")
-
-# Define the template for parsing
 template = (
     "You are tasked with extracting specific information from the following text content: {dom_content}. "
     "Please follow these instructions carefully: \n\n"
@@ -13,23 +10,19 @@ template = (
     "4. **Direct Data Only:** Your output should contain only the data that is explicitly requested, with no other text."
 )
 
-def parse_with_groq(dom_chunks, parse_description, model="llama3-8b-8192"):
+model = OllamaLLM(model="llama3")
+
+def parse_with_ollama(dom_chunks, parse_description):
+    prompt = ChatPromptTemplate.from_template(template)
+    chain = prompt | model
+
     parsed_results = []
 
     for i, chunk in enumerate(dom_chunks, start=1):
-        # Prepare the prompt
-        prompt = template.format(dom_content=chunk, parse_description=parse_description)
-
-        # Send prompt to Groq for processing, specifying the model
-        response = client.chat.completions.create(
-            messages=[
-                {"role": "user", "content": prompt}
-            ],
-            model=model  # Specify the model
+        response = chain.invoke(
+            {"dom_content": chunk, "parse_description": parse_description}
         )
-        
-        # Print status and store result
         print(f"Parsed batch: {i} of {len(dom_chunks)}")
-        parsed_results.append(response.choices[0].message.content)  # Access the content
+        parsed_results.append(response)
 
     return "\n".join(parsed_results)