Spaces:

abhibisht89
/

neural-search-engine

Running

abhibisht89

Fix Space runtime compatibility

a4fa5f9 4 months ago

2.95 kB

import pickle
from pathlib import Path

import gradio as gr
from sentence_transformers import CrossEncoder, SentenceTransformer, util

DATA_DIR = Path(__file__).parent
RESULT_COUNT = 5


def load_pickle(filename):
    """Load one of the trusted data files stored with this Space."""
    # Pickle files can execute code while loading. Only load the two files
    # committed to this repository. Never pass user uploads to this function.
    with (DATA_DIR / filename).open("rb") as file:
        return pickle.load(file)


# Load each model once when the Space starts. The corrected reranker ID matches
# the official Sentence Transformers model repository.
bi_encoder = SentenceTransformer("multi-qa-MiniLM-L6-cos-v1")
cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2")
corpus_embeddings = load_pickle("corpus_embeddings_cpu.pkl")
corpus = load_pickle("corpus.pkl")


def search(query, top_k=100):
    """Retrieve relevant passages, rerank them, and return five results."""
    query = (query or "").strip()
    if not query:
        return ("Enter a search query.", "", "", "", "")

    # Semantic search quickly finds a small set of candidate passages.
    # Encode the query using the bi-encoder and find potentially relevant passages
    question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
    hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k)
    hits = hits[0]  # Get the hits for the first query

    # The slower cross-encoder improves result quality by reranking candidates.
    # Now, score all retrieved passages with the cross_encoder
    cross_inp = [[query, corpus[hit['corpus_id']]] for hit in hits]
    cross_scores = cross_encoder.predict(cross_inp)

    # Sort results by the cross-encoder scores
    for hit, score in zip(hits, cross_scores):
        hit["cross_score"] = float(score)

    hits = sorted(hits, key=lambda hit: hit["cross_score"], reverse=True)
    answers = [corpus[hit["corpus_id"]] for hit in hits[:RESULT_COUNT]]

    # Gradio expects one value for each of the five output boxes.
    answers.extend([""] * (RESULT_COUNT - len(answers)))
    return tuple(answers)


examples = [
    "Who is Steve Jobs?",
    "What is Coldplay?",
    "What is a Turing test?",
    "What is the most interesting thing about our universe?",
    "What are the most beautiful places on Earth?",
]

description = (
    "This semantic search engine uses Sentence Transformers to retrieve and "
    "rerank Wikipedia passages. It returns the five most relevant results."
)

search_input = gr.Textbox(lines=1, label="Search query")
outputs = [
    gr.Textbox(type="text", label=f"Search result {index}")
    for index in range(1, RESULT_COUNT + 1)
]

interface = gr.Interface(
    fn=search,
    inputs=search_input,
    outputs=outputs,
    examples=examples,
    article=description,
    title="Neural Search Engine",
)

if __name__ == "__main__":
    interface.launch()

	import pickle
	from pathlib import Path

	import gradio as gr
	from sentence_transformers import CrossEncoder, SentenceTransformer, util

	DATA_DIR = Path(__file__).parent
	RESULT_COUNT = 5


	def load_pickle(filename):
	"""Load one of the trusted data files stored with this Space."""
	# Pickle files can execute code while loading. Only load the two files
	# committed to this repository. Never pass user uploads to this function.
	with (DATA_DIR / filename).open("rb") as file:
	return pickle.load(file)


	# Load each model once when the Space starts. The corrected reranker ID matches
	# the official Sentence Transformers model repository.
	bi_encoder = SentenceTransformer("multi-qa-MiniLM-L6-cos-v1")
	cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2")
	corpus_embeddings = load_pickle("corpus_embeddings_cpu.pkl")
	corpus = load_pickle("corpus.pkl")


	def search(query, top_k=100):
	"""Retrieve relevant passages, rerank them, and return five results."""
	query = (query or "").strip()
	if not query:
	return ("Enter a search query.", "", "", "", "")

	# Semantic search quickly finds a small set of candidate passages.
	# Encode the query using the bi-encoder and find potentially relevant passages
	question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
	hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k)
	hits = hits[0] # Get the hits for the first query

	# The slower cross-encoder improves result quality by reranking candidates.
	# Now, score all retrieved passages with the cross_encoder
	cross_inp = [[query, corpus[hit['corpus_id']]] for hit in hits]
	cross_scores = cross_encoder.predict(cross_inp)

	# Sort results by the cross-encoder scores
	for hit, score in zip(hits, cross_scores):
	hit["cross_score"] = float(score)

	hits = sorted(hits, key=lambda hit: hit["cross_score"], reverse=True)
	answers = [corpus[hit["corpus_id"]] for hit in hits[:RESULT_COUNT]]

	# Gradio expects one value for each of the five output boxes.
	answers.extend([""] * (RESULT_COUNT - len(answers)))
	return tuple(answers)


	examples = [
	"Who is Steve Jobs?",
	"What is Coldplay?",
	"What is a Turing test?",
	"What is the most interesting thing about our universe?",
	"What are the most beautiful places on Earth?",
	]

	description = (
	"This semantic search engine uses Sentence Transformers to retrieve and "
	"rerank Wikipedia passages. It returns the five most relevant results."
	)

	search_input = gr.Textbox(lines=1, label="Search query")
	outputs = [
	gr.Textbox(type="text", label=f"Search result {index}")
	for index in range(1, RESULT_COUNT + 1)
	]

	interface = gr.Interface(
	fn=search,
	inputs=search_input,
	outputs=outputs,
	examples=examples,
	article=description,
	title="Neural Search Engine",
	)

	if __name__ == "__main__":
	interface.launch()