File size: 2,732 Bytes
e64b928 bcee4f0 e64b928 e9273c1 e64b928 6cc2eec e64b928 6cc2eec e64b928 6cc2eec e64b928 6cc2eec e64b928 6cc2eec e64b928 e9273c1 e64b928 6cc2eec e64b928 9e0468e e64b928 6cc2eec e64b928 9e0468e e64b928 9e0468e e64b928 6cc2eec e64b928 e9273c1 4d69305 e9273c1 e64b928 e9273c1 e64b928 4986b8a e64b928 6cc2eec e64b928 3b9de69 7635d6b 3b9de69 e64b928 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | 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() |