parse.py

import streamlit as st
from scraper import scrape_website, split_dom_content, clean_body_content, extract_body_content
from parse import parse, merge_tables_with_llm
from Data import markdown_to_csv
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
import pandas as pd

# Initialiser le modèle LLM avec transformers
model_name = "meta-llama/Llama-4-Maverick-17B-128E-Instruct"  # Modèle léger pour cet exemple
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Créer un pipeline pour la génération de texte
llm_pipeline = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device=0 if torch.cuda.is_available() else -1,  # Utiliser GPU si disponible
    max_new_tokens=500,  # Limiter la longueur de la réponse
)

st.title("AI Web Scraper")

# Multi-URL Input
urls = st.text_area("Enter Website URLs (one per line)", height=150)
urls_list = [url.strip() for url in urls.splitlines() if url.strip()]

if st.button("Scrape Sites"):
    all_results = []
    for url in urls_list:
        st.write(f"Scraping: {url}")
        result = scrape_website(url)
        body_content = extract_body_content(result)
        cleaned_content = clean_body_content(body_content)
        all_results.append(cleaned_content)

    st.session_state.all_dom_content = all_results

if "all_dom_content" in st.session_state:
    parse_description = st.text_area("Describe what you want to parse from ALL sites:")

    if st.button("Parse Content"):
        if parse_description:
            all_tables = []
            for i, dom_content in enumerate(st.session_state.all_dom_content):
                st.write(f"Parsing content from site {i+1}")
                dom_chunks = split_dom_content(dom_content)
                result = parse(dom_chunks, parse_description)  # Assurez-vous que cette fonction utilise le modèle local
                st.write("Raw LLM Output:")
                st.write(result)

                tables = markdown_to_csv(result)
                if tables:
                    st.write("Extracted Tables:")
                    for table in tables:
                        st.write(table)
                        all_tables.append(table)
                else:
                    st.write("No tables found in the output. Displaying raw output instead.")
                    st.text_area("Raw Output", result, height=200)

            # Merge tables using LLM
            if all_tables:
                st.write("Merging all tables using LLM...")
                merged_table_string = merge_tables_with_llm(all_tables, parse_description)
                st.write("Merged Table (LLM Output):")
                st.write(merged_table_string)

                # Convert merged table string to DataFrame
                merged_tables = markdown_to_csv(merged_table_string)
                if merged_tables:
                    st.write("Merged Table (DataFrame):")
                    st.write(merged_tables[0])
                else:
                    st.write("Could not convert merged table string to DataFrame.")
            else:
                st.write("No tables to merge.")

def merge_tables_with_llm(tables, parse_description):
    """Merges a list of Pandas DataFrames into a single Markdown table using a local LLM."""
    # Convert DataFrames to Markdown strings
    table_strings = [table.to_markdown(index=False) for table in tables]

    # Créer un prompt pour le LLM
    merge_prompt = (
        "You are tasked with merging the following Markdown tables into a single, comprehensive Markdown table.\n"
        "The tables contain information related to: {parse_description}.\n"
        "Combine the tables, ensuring that the merged table is well-formatted and contains all relevant information.\n"
        "If there are duplicate columns, rename them to be unique. If there are missing values, fill them with 'N/A'.\n"
        "Ensure the final output is a single valid Markdown table.\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)

    # Appeler le modèle local via le pipeline
    response = llm_pipeline(merge_prompt, max_length=2000, truncation=True)
    merged_table = response[0]["generated_text"]

    # Nettoyer la sortie pour ne garder que la table Markdown
    # Supposons que la table commence après le prompt
    start_idx = merged_table.find("|")
    if start_idx != -1:
        merged_table = merged_table[start_idx:]
    else:
        merged_table = "No valid Markdown table found in LLM output."

    return merged_table