app.py

import streamlit as st
import torch
import io
import re
import numpy as np
import pandas as pd
from threading import Thread

# LLM
from transformers.models.auto.tokenization_auto import AutoTokenizer
from transformers.models.auto.modeling_auto import AutoModelForCausalLM
from transformers.generation.streamers import TextIteratorStreamer

# Embedding + Reranker
from langchain_huggingface import HuggingFaceEmbeddings
from sentence_transformers import CrossEncoder

# Document processing
from langchain_community.vectorstores import FAISS
from pypdf import PdfReader

# BM25 (keyword search)
from rank_bm25 import BM25Okapi

# ==========================================
# ⚡ CPU CONFIG — free HuggingFace Spaces
# ==========================================
# ⚡ CPU CONFIG — call this before any parallel work
try:
    torch.set_num_threads(2)
    torch.set_num_interop_threads(1)
except RuntimeError:
    pass  # Already initialized by Streamlit at startup; ignore.


    
# Layout 

st.set_page_config(page_title="Hybrid RAG + Reranker", page_icon="🔍", layout="wide")
st.title("🔍 Hybrid RAG + Reranker")

# Sidebar text 

with st.sidebar.expander("⚙️ How does this system work?", expanded=False):
    st.markdown("""
**Step 1 - Dual retrieval (hybrid)**

When you ask a question, two search engines work in parallel over your chunks:

- **BM25**: classic keyword search (like Google in the 1990s). Strong on names, dates, and exact terms. Very fast, no AI model.
- **FAISS**: semantic vector search. Understands meaning even when the words differ.

Each returns its own **top 5 chunks** -> we merge them -> **10 candidate chunks**.

---

**Step 2 - Reranking (cross-encoder)**

A small model (~70M params) reads each *(question, chunk)* pair together and assigns a true relevance score.

Unlike embeddings that encode the question and chunk **separately**, the cross-encoder reads them **together** -> it actually understands the relationship.

Result: we keep the **2 best chunks** out of the 10.

---

**Step 3 - Generation**

The LLM receives only these 2 highly relevant chunks and generates the answer. One LLM call, clean context = precise and fast response.
    """)

# Load the models once (embedding + reranker + LLM + quantization)
@st.cache_resource(show_spinner="Loading models... (first run ~2 min)")
def load_models():

    # ⚡ Embedding: MiniLM 120M - fast and sufficient
    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
        model_kwargs={'device': 'cpu'},
        encode_kwargs={'normalize_embeddings': True, 'batch_size': 16}
    )

    # ⚡ Reranker: lightweight cross-encoder, ~70M params
    # Reads (question + chunk) together -> real relevance score
    # Much more accurate than simple cosine similarity
    reranker = CrossEncoder(
        "cross-encoder/ms-marco-MiniLM-L-6-v2",
        max_length=512,
        device="cpu"
    )

    # LLM: Qwen 1.5B Instruct
    model_id = "Qwen/Qwen2.5-0.5B-Instruct"  # instead of 1.5B
    tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
    

    llm = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype=torch.bfloat16,  # instead of float32
        low_cpu_mem_usage=True
    )
    
    # Warm-up: JIT-compile the kernels and avoids lag on the first request
    dummy = tokenizer("warmup", return_tensors="pt")
    with torch.inference_mode():
        llm.generate(**dummy, max_new_tokens=1, pad_token_id=tokenizer.eos_token_id)

    return embeddings, reranker, tokenizer, llm

embeddings, reranker, tokenizer, llm = load_models()

# Qwen stop tokens
EOS_IDS = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|im_end|>")]
MAX_TOKENS_ANSWER = 384

# ==========================================
# SESSION STATE
# ==========================================
for key, default in [
    ("vector_db", None),
    ("bm25", None),
    ("chunks_text", None),
    ("current_file_name", None)
]:
    if key not in st.session_state:
        st.session_state[key] = default

# PDF upload section (allow multiple files)
uploaded_files = st.sidebar.file_uploader("📂 Upload your PDFs (multiple allowed)", type=["pdf"], accept_multiple_files=True)

if not uploaded_files:
    for key in ["vector_db", "bm25", "chunks_text", "current_file_names", "chunks_source"]:
        st.session_state[key] = None

else:
    # Collect filenames
    filenames = [f.name for f in uploaded_files]

    # Only re-index if uploaded set changed
    if st.session_state.get("current_file_names") != filenames:
        for key in ["vector_db", "bm25", "chunks_text", "chunks_source"]:
            st.session_state[key] = None
        st.session_state.current_file_names = filenames

    if st.session_state.vector_db is None:
        with st.spinner("Indexing documents... ⏳"):
            all_chunks = []
            all_sources = []

            # Extract and chunk each uploaded PDF
            for f in uploaded_files:
                try:
                    reader = PdfReader(io.BytesIO(f.getvalue()))
                    full_text = "\n".join(page.extract_text() or "" for page in reader.pages)
                    normalized_text = re.sub(r"\s+", " ", full_text).strip()

                    def split_text_into_chunks(text: str, chunk_size: int = 500, chunk_overlap: int = 200) -> list[str]:
                        if not text:
                            return []

                        chunks = []
                        start = 0
                        text_length = len(text)
                        while start < text_length:
                            end = min(start + chunk_size, text_length)
                            chunks.append(text[start:end].strip())
                            if end >= text_length:
                                break
                            start = max(end - chunk_overlap, 0)
                        return [chunk for chunk in chunks if chunk]

                    chunks = split_text_into_chunks(normalized_text)
                    all_chunks.extend(chunks)
                    all_sources.extend([f.name] * len(chunks))
                except Exception as e:
                    st.sidebar.warning(f"Could not process {f.name}: {e}")

            # --- FAISS index (semantic search) ---
            st.session_state.vector_db = FAISS.from_texts(all_chunks, embeddings)

            # --- BM25 index (keyword search) ---
            tokenized = [t.lower().split() for t in all_chunks]
            st.session_state.bm25 = BM25Okapi(tokenized)
            st.session_state.chunks_text = all_chunks   # raw text for BM25
            st.session_state.chunks_source = all_sources

            st.sidebar.success(f"✅ Indexed {len(all_chunks)} chunks from {len(filenames)} files")


# Sidebar: scope selector (All files or specific uploaded file)
scope_options = ["All files"]
if st.session_state.get("current_file_names"):
    scope_options += st.session_state.current_file_names
selected_scope = st.sidebar.selectbox("Index scope", scope_options, index=0)
selected_file = None if selected_scope == "All files" else selected_scope




# Fonctions

def hybrid_search(question: str, k: int = 10, source_filter: str | None = None) -> list[str]:
    """
    Hybrid search: BM25 + FAISS -> merge the results.

    Why k=5 for each?
    We retrieve broadly on purpose (10 candidates total)
    because the reranker will sort and keep the true best matches.
    BM25 catches what FAISS misses (exact terms) and vice versa.
    """
    question_lower = question.lower()

    # -- BM25: term-frequency score --
    bm25_scores = st.session_state.bm25.get_scores(question_lower.split())
    top_bm25_idx = np.argsort(bm25_scores)[::-1][:k]
    bm25_results = []
    for i in top_bm25_idx:
        text = st.session_state.chunks_text[i]
        if source_filter is None or st.session_state.chunks_source[i] == source_filter:
            bm25_results.append(text)

    # -- FAISS: semantic similarity score --
    faiss_docs   = st.session_state.vector_db.similarity_search(question, k=k)
    faiss_results = []
    for d in faiss_docs:
        # d may be a Document-like object
        text = getattr(d, 'page_content', str(d))
        # try to find index to check source
        try:
            idx = st.session_state.chunks_text.index(text)
            source = st.session_state.chunks_source[idx]
        except Exception:
            source = None
        if source_filter is None or source == source_filter:
            faiss_results.append(text)

    # -- Merge with deduplication --
    seen  = set()
    merged = []
    for text in bm25_results + faiss_results:
        if text not in seen:
            seen.add(text)
            merged.append(text)

    return merged  # up to 10 unique chunks


def rerank(question: str, candidates: list[str], top_n: int = 2):
    pairs  = [(question, chunk) for chunk in candidates]
    scores = reranker.predict(pairs)
    ranked = sorted(zip(scores, candidates), reverse=True)
    best_score = ranked[0][0]
    best_chunks = [chunk for _, chunk in ranked[:top_n]]
    return best_chunks, float(best_score)


def build_prompt(system: str, user: str) -> str:
    return (
        f"<|im_start|>system\n{system}\n<|im_end|>\n"
        f"<|im_start|>user\n{user}\n<|im_end|>\n"
        f"<|im_start|>assistant\n"
    )


def stream_llm(prompt: str, placeholder) -> str:
    """Streamed generation with an animated cursor."""
    inputs   = tokenizer(prompt, return_tensors="pt")
    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=True
    )

    def _generate():
        with torch.inference_mode():
            llm.generate(
                **inputs,
                streamer=streamer,
                max_new_tokens=MAX_TOKENS_ANSWER,
                do_sample=False,
                repetition_penalty=1.3,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=EOS_IDS
            )

    Thread(target=_generate).start()

    full_text = ""
    for token in streamer:
        full_text += token
        placeholder.markdown(full_text + "▌")
    placeholder.markdown(full_text)
    return full_text

def clean_chunk(text: str) -> str:
    """Mark truncated chunks instead of sending them raw."""
    text = text.strip()
    # If the chunk starts lowercase, it is likely truncated.
    if text and text[0].islower():
        text = "(...) " + text
    return text


def classify_question_style(question: str) -> str:
    question_lower = question.lower()

    value_markers = [
        "value", "values", "calculate", "calculation", "compute", "how many", "how much",
        "what is the amount", "report the number", "report the value", "percentage", "percent",
        "%", "ratio", "emissions", "ghg", "co2", "scope 1", "scope 2", "scope 3", "tonnes", "tons"
    ]
    structured_markers = [
        "list", "provide", "state", "report", "identify", "name", "names", "describe",
        "what are", "which", "give the", "enumerate", "include"
    ]

    if any(marker in question_lower for marker in value_markers):
        return "value_only"
    if any(marker in question_lower for marker in structured_markers):
        return "structured"
    return "paragraph"


def build_answer_system_prompt(answer_style: str) -> str:
    if answer_style == "value_only":
        return (
            "You are a precise extraction engine. "
            "Return only the exact final value, number, percentage, or unit requested by the question. "
            "Do not add explanation, labels, or narrative text. "
            "If the answer requires a calculation, compute it from the context and return the final value only. "
            "If the answer is not present, say: Information not found."
        )

    if answer_style == "structured":
        return (
            "You are a structured extractor. "
            "Return a concise answer using short sentences or bullet points, preserving all requested facts. "
            "Do not compress multiple requested fields into one vague paragraph. "
            "If the question asks for several items, answer each item explicitly. "
            "Do not omit numbers, names, dates, locations, or amounts that are relevant. "
            "If the answer is not present, say: Information not found."
        )

    return (
        "You are a fact extractor. "
        "Use the context to answer fully and faithfully, without compressing away important facts. "
        "If the answer contains multiple facts, include them all in a clear paragraph. "
        "Do not omit numbers, names, dates, locations, or amounts that are relevant. "
        "Never generate URLs, emails, menus, or language lists. "
        "If the context contains '(...)', it is an incomplete sentence, so reconstruct the meaning. "
        "If you cannot find the answer, say: Information not found."
    )


def build_answer_instruction(answer_style: str) -> str:
    if answer_style == "value_only":
        return (
            "Answer format: output only the final value, number, percentage, or unit. "
            "Do not explain your reasoning."
        )
    if answer_style == "structured":
        return (
            "Answer format: use short structured points or a compact list, one item per requested field."
        )
    return (
        "Answer format: provide a full, faithful paragraph with all relevant facts."
    )


def generate_answer(question: str, context: str, batch_mode: bool = False) -> str:
    answer_style = classify_question_style(question)
    prompt = build_prompt(
        system=build_answer_system_prompt(answer_style),
        user=(
            f"CONTEXT:\n{context}\n\n"
            f"QUESTION: {question}\n\n"
            f"{build_answer_instruction(answer_style)}"
        )
    )

    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.inference_mode():
        output = llm.generate(
            **inputs,
            max_new_tokens=MAX_TOKENS_ANSWER if batch_mode else MAX_TOKENS_ANSWER,
            do_sample=False,
            repetition_penalty=1.05 if answer_style == "value_only" else 1.15,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=EOS_IDS
        )

    generated_tokens = output[0][inputs["input_ids"].shape[-1]:]
    answer_text = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
    return answer_text

# ==========================================
# MAIN INTERFACE
# ==========================================
if st.session_state.vector_db is not None:

    question = st.text_input(
        "Ask your question:",
        placeholder="Type a question"
    )

    if st.button("Run analysis", type="primary") and question:
        st.markdown("---")

        # --- STEP 1: HYBRID RETRIEVAL ---
        with st.spinner("🔎 Hybrid retrieval (BM25 + vector search)..."):
            candidates = hybrid_search(question, k=5, source_filter=selected_file)

        # Optional debug display
        with st.expander(f"📚 {len(candidates)} candidate chunks retrieved (before reranking)"):
            for i, c in enumerate(candidates):
                st.caption(f"**Chunk {i+1}**: {c[:150]}...")

        # --- STEP 2: RERANKING ---
        with st.spinner("⚖️ Reranking (selecting the 2 best chunks)..."):
            best_chunks, confidence = rerank(question, candidates, top_n=2)

        st.caption(f"Relevance score: `{confidence:.2f}`")
        
        if confidence < -2.0:  # very negative score = nothing relevant
            st.warning("❌ No relevant passage found for this question.")
            st.stop()

        with st.expander("✅ 2 chunks kept after reranking"):
            for i, c in enumerate(best_chunks):
                st.info(f"**Chunk {i+1}**: {c}")

        # --- STEP 3: LLM GENERATION ---
        context = "\n\n".join([clean_chunk(c) for c in best_chunks])

        st.markdown("### 🎯 Answer:")
        placeholder = st.empty()
        placeholder.markdown(generate_answer(question, context))

    # ==========================================
    # INDICATORS & BATCH ANALYSIS
    # ==========================================
    st.markdown("---")
    st.markdown("## 📋 Batch Analysis: All Indicators")

    @st.cache_resource(show_spinner=False)
    def load_indicators():
        """Parse the indicators and questions from the file."""
        indicators = []
        try:
            with open("all_indicators_with_questions - Copy.txt", "r", encoding="utf-8") as f:
                lines = f.readlines()
            
            i = 0
            while i < len(lines):
                line = lines[i].strip()
                
                # Match indicator lines (e.g., "1. IMP-M01-I01 | ...")
                if line and line[0].isdigit() and "|" in line:
                    parts = line.split("|")
                    if len(parts) >= 3:
                        id_part = parts[0].strip().split(".")[-1].strip()
                        category = parts[1].strip()
                        indicator_name = parts[2].strip()
                        
                        # Next line should be the question
                        i += 1
                        if i < len(lines) and lines[i].strip().startswith("Question:"):
                            question_text = lines[i].strip()[10:].strip()
                            indicators.append({
                                "id": id_part,
                                "category": category,
                                "indicator": indicator_name,
                                "question": question_text,
                                "answer": ""
                            })
                i += 1
        except Exception as e:
            st.error(f"Error loading indicators: {str(e)}")
        
        return indicators

    indicators = load_indicators()

    if indicators:
        col1, col2 = st.columns([3, 1])
        with col1:
            st.write(f"**Found {len(indicators)} indicators**")
        with col2:
            run_batch_analysis = st.button("🚀 Analyze All", type="primary", use_container_width=True)

        if run_batch_analysis:
            # Create a progress bar and results storage
            progress_bar = st.progress(0)
            status_text = st.empty()
            results_placeholder = st.empty()

            results = []

            for idx, ind in enumerate(indicators):
                # Update progress
                progress = (idx) / len(indicators)
                progress_bar.progress(progress)
                status_text.write(f"Processing {idx + 1}/{len(indicators)}: **{ind['id']}** - {ind['indicator'][:50]}...")

                # Run the analysis for this indicator
                question = ind["question"]
                
                try:
                    # Hybrid search
                    candidates = hybrid_search(question, k=5, source_filter=selected_file)
                    
                    if candidates:
                        # Reranking
                        best_chunks, confidence = rerank(question, candidates, top_n=2)
                        
                        # Generate answer
                        if confidence >= -2.0:
                            context = "\n\n".join([clean_chunk(c) for c in best_chunks])
                            answer_text = generate_answer(question, context, batch_mode=True)
                            
                            results.append({
                                **ind,
                                "answer": answer_text,
                                "confidence": f"{confidence:.2f}"
                            })
                        else:
                            results.append({**ind, "answer": "No relevant passage found.", "confidence": f"{confidence:.2f}"})
                    else:
                        results.append({**ind, "answer": "No relevant passages retrieved.", "confidence": "N/A"})
                        
                except Exception as e:
                    results.append({**ind, "answer": f"Error: {str(e)[:50]}", "confidence": "N/A"})

            # Complete progress
            progress_bar.progress(1.0)
            status_text.write(f"✅ Completed: {len(results)} indicators analyzed")

            # Display results in a table
            if results:
                st.subheader("📊 Analysis Results")
                
                # Create a DataFrame for display
                df_results = pd.DataFrame([
                    {
                        "ID": r["id"],
                        "Category": r["category"],
                        "Indicator": r["indicator"],
                        "Question": r["question"],
                        "Answer": r["answer"],
                        "Score": r["confidence"]
                    }
                    for r in results
                ])

                df_display = df_results.copy()
                df_display["Indicator"] = df_display["Indicator"].apply(lambda value: value[:40] + "..." if len(value) > 40 else value)
                df_display["Question"] = df_display["Question"].apply(lambda value: value[:60] + "..." if len(value) > 60 else value)
                df_display["Answer"] = df_display["Answer"].apply(lambda value: value[:120] + "..." if len(value) > 120 else value)
                
                st.dataframe(df_display, use_container_width=True, height=500)
                
                # Option to export results
                csv_output = df_results.to_csv(index=False)
                st.download_button(
                    label="📥 Download Results as CSV",
                    data=csv_output,
                    file_name="indicators_analysis_results.csv",
                    mime="text/csv"
                )
    else:
        st.warning("⚠️ Could not load indicators. Make sure 'all_indicators_with_questions - Copy.txt' is in the current directory.")

else:
    st.info("⬅️ Upload a PDF file in the sidebar to begin.")