| 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]) |
|
|
| |
| 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)) |
|
|
| |
| model_name = "BAAI/bge-small-en-v1.5" |
| tokenizer = AutoTokenizer.from_pretrained(model_name) |
| emb_model = AutoModel.from_pretrained(model_name) |
|
|
| |
| 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) |
|
|
| |
| 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 |
| ) |
| |
| |
| 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() |
