import gradio as gr
import faiss
import numpy as np
from transformers import AutoTokenizer, AutoModel
import pickle
import torch

from sentence_transformers import CrossEncoder


index = faiss.read_index("index503.faiss")

with open("texts503.pkl", "rb") as f:
    texts = pickle.load(f)

print(texts[500])

# chunking
chunks = []
chunk_size = 25
overlap_int = 2
for text in texts:
    text_words = text.split()
    for i in range(0, len(text_words), chunk_size):
        chunk_words = text_words[max(i - overlap_int, 0): i + chunk_size]
        chunk = " ".join(chunk_words)
        chunks.append(chunk)

print(len(chunks))

# emb model
model_name = "BAAI/bge-small-en-v1.5"
tokenizer = AutoTokenizer.from_pretrained(model_name)
emb_model = AutoModel.from_pretrained(model_name)

#rerank model
reranker = CrossEncoder(
    "cross-encoder/ms-marco-MiniLM-L2-v2",
    max_length=512,
    device="cpu"
)

def embed(query):

    inputs = tokenizer(
        query,
        padding=True,
        truncation=True,
        return_tensors="pt"
    )

    with torch.no_grad():
        outputs = emb_model(**inputs)

    # Mean pooling: match your PC's implementation
    embeddings = outputs.last_hidden_state.mean(dim=1)

    return embeddings.cpu().numpy().astype("float32")


def retrieve(query, k=5):

    q = embed([query])
    scores, ids = index.search(q, k)
    
    retrieved_chunks = [chunks[i] for i in ids[0]]

    first_ret = "\n".join("Chunk ID"+str(i) +": " + chunks[i] for i in ids[0])
    

    pairs = [
        (query, chunk)
        for chunk in retrieved_chunks
    ]
    
    rerank_scores = reranker.predict(
        pairs,
        batch_size=4
    )
    
    ranked_chunks = sorted(
        zip(retrieved_chunks, rerank_scores),
        key=lambda x: x[1],
        reverse=True
    )
    
    # 3. Format reranked results
    reranked_ret = "\n".join(
        f"{i}: {chunk} (score: {score:.4f})"
        for i, (chunk, score) in enumerate(ranked_chunks, 1)
    )

    result = first_ret + "\n\nAfter reranking:\n" + reranked_ret

    return result
    


with gr.Blocks() as demo:
    gr.Markdown("""# Enter anything and press Enter""")
    
    topic_input = gr.Textbox(
        label="Enter here"
    )

    output = gr.Textbox(
        label="Top5 most relevant news:"
    )

    topic_input.submit(
        retrieve,
        inputs=topic_input,
        outputs=output
    )

    gr.Markdown("""
- **Dataset: 500 News from AG_News Data + 3 custom News start with Custom News by Curry MindFree**
- **Chunking size: 25 words. Overlap: 2 words**
- **Vectorize model: BAAI/bge-small-en-v1.5**
- **Vectors similarity search: FAISS Indexing**
- **Rerank model: cross-encoder/ms-marco-MiniLM-L2-v2**
""")

demo.launch()