Update parse.py

parse.py

@@ -1,93 +1,59 @@
-import pandas as pd
 from langchain_core.prompts import ChatPromptTemplate
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
-from huggingface_hub import login
-import torch
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage
 import os
+import pandas as pd
 
-# Hugging Face API Token from Space Secrets
-HF_TOKEN = os.getenv("HF_TOKEN")
-if not HF_TOKEN:
-    raise ValueError("HF_TOKEN environment variable not set. Please set it in Hugging Face Space Settings under Secrets.")
-
-# Model configuration
-MODEL_NAME = "mistralai/Mixtral-8x7B-Instruct-v0.1"  
-
-# Initialize model and tokenizer
-try:
-    # Log in to Hugging Face Hub
-    login(token=HF_TOKEN)
-    print("Successfully logged in to Hugging Face Hub")
 
-    # Load tokenizer and model
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
-    # Create text generation pipeline
-    llm_pipeline = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        device=0 if torch.cuda.is_available() else -1,  # Use GPU if available in Space
-        max_new_tokens=500,  # Limit response length
-        pad_token_id=tokenizer.eos_token_id,  # Ensure proper padding
-    )
-except Exception as e:
-    print(f"Failed to load model: {str(e)}")
-    llm_pipeline = None
+# Load OpenRouter API Key
+openrouter_api_key = os.getenv("API_MV")
+
+model = ChatOpenAI(
+    openai_api_key=openrouter_api_key,  # Use OpenRouter API key
+    model="meta-llama/llama-4-maverick:free",  # Specify Qwen VL Plus model
+    base_url="https://openrouter.ai/api/v1"  # OpenRouter API URL
+)
+
+# Create a chat prompt template
+template = (
+    "You are tasked with extracting specific information from the following text content: {dom_content}. "
+    "Please follow these instructions carefully:\n\n"
+    "1. **Task:** Extract data from the provided text that matches the description: {parse_description}.\n"
+    "2. **Output Format:** Return the extracted data ONLY as one or more Markdown tables. Each table MUST be correctly formatted.\n"
+    "3. **Markdown Table Format:** Each table must adhere to the following Markdown format:\n"
+    "   - Start with a header row, clearly labeling each column, separated by pipes (|).\n"
+    "   - Follow the header row with an alignment row, using hyphens (-) to indicate column alignment (e.g., --- for left alignment).\n"
+    "   - Subsequent rows should contain the data, with cells aligned according to the alignment row.\n"
+    "   - Use pipes (|) to separate columns in each data row.\n"
+    "4. **No Explanations:** Do not include any introductory or explanatory text before or after the table(s).\n"
+    "5. **Empty Response:** If no information matches the description, return an empty string ('').\n"
+    "6. **Multiple Tables:** If the text contains multiple tables matching the description, return each table separately, following the Markdown format for each.\n"
+    "7. **Accuracy:** The extracted data must be accurate and reflect the information in the provided text.\n"
+)
 
 # Function to parse and extract information from the chunks
 def parse(dom_chunks, parse_description):
-    """Parse and extract information from DOM chunks using a local LLM."""
-    if llm_pipeline is None:
-        raise ValueError("LLM pipeline not initialized. Check model loading and ensure HF_TOKEN is set in Space Secrets.")
-
-    # Create a prompt template
-    template = (
-        "You are tasked with extracting specific information from the following text content: {dom_content}. "
-        "Please follow these instructions carefully:\n\n"
-        "1. **Task:** Extract data from the provided text that matches the description: {parse_description}.\n"
-        "2. **Output Format:** Return the extracted data ONLY as one or more Markdown tables. Each table MUST be correctly formatted.\n"
-        "3. **Markdown Table Format:** Each table must adhere to the following Markdown format:\n"
-        "   - Start with a header row, clearly labeling each column, separated by pipes (|).\n"
-        "   - Follow the header row with an alignment row, using hyphens (-) to indicate column alignment (e.g., --- for left alignment).\n"
-        "   - Subsequent rows should contain the data, with cells aligned according to the alignment row.\n"
-        "   - Use pipes (|) to separate columns in each data row.\n"
-        "4. **No Explanations:** Do not include any introductory or explanatory text before or after the table(s).\n"
-        "5. **Empty Response:** If no information matches the description, return an empty string ('').\n"
-        "6. **Multiple Tables:** If the text contains multiple tables matching the description, return each table separately, following the Markdown format for each.\n"
-        "7. **Accuracy:** The extracted data must be accurate and reflect the information in the provided text.\n"
-    )
+    prompt = ChatPromptTemplate.from_template(template)
+    chain = prompt | model
 
     parsed_results = []
 
     # Loop through the chunks and parse
     for i, chunk in enumerate(dom_chunks, start=1):
-        # Format the prompt
-        prompt = template.format(dom_content=chunk, parse_description=parse_description)
-        
-        # Invoke the LLM pipeline
-        response = llm_pipeline(prompt, max_length=2000, truncation=True)
-        result = response[0]["generated_text"]
-        
-        # Clean the output to keep only the Markdown table (remove prompt text)
-        start_idx = result.find("|")
-        if start_idx != -1:
-            result = result[start_idx:]
-        else:
-            result = ""  # Return empty string if no table is found
+        response = chain.invoke({"dom_content": chunk, "parse_description": parse_description})
 
+        # Extract the content from AIMessage and add it to the results
         print(f"Parsed batch {i} of {len(dom_chunks)}")
-        parsed_results.append(result)
+        parsed_results.append(response.content)  # Ensure content is extracted properly
 
     # Return the parsed results as a single string
     return "\n".join(parsed_results)
 
 def merge_tables_with_llm(tables, parse_description):
-    """Merges a list of Pandas DataFrames into a single Markdown table using a local LLM."""
-    if llm_pipeline is None:
-        raise ValueError("LLM pipeline not initialized. Check model loading and ensure HF_TOKEN is set in Space Secrets.")
-
+    """Merges a list of Pandas DataFrames into a single Markdown table using LLM."""
+    from langchain_core.prompts import ChatPromptTemplate
+    from langchain_openai import ChatOpenAI
     # Convert DataFrames to Markdown strings
     table_strings = [table.to_markdown(index=False) for table in tables]
 
@@ -109,17 +75,9 @@ def merge_tables_with_llm(tables, parse_description):
         "7. **Missing Values:** If there are missing values, fill them with 'N/A'.\n\n"
         "Here are the tables:\n\n" + "\n\n".join(table_strings) +
         "\n\nReturn the merged table in Markdown format:"
-    ).format(parse_description=parse_description)
-
-    # Invoke the LLM pipeline
-    response = llm_pipeline(merge_prompt, max_length=2000, truncation=True)
-    merged_table = response[0]["generated_text"]
-
-    # Clean the output to keep only the Markdown table
-    start_idx = merged_table.find("|")
-    if start_idx != -1:
-        merged_table = merged_table[start_idx:]
-    else:
-        merged_table = ""  # Return empty string if no table is found
+    )
 
-    return merged_table
+    # Invoke the LLM
+    message = HumanMessage(content=merge_prompt)
+    response = model.invoke([message])
+    return response.content