import argparse import logging import datasets import gradio as gr import sentence_transformers logging.disable(logging.CRITICAL) model = sentence_transformers.SentenceTransformer( "dangvantuan/sentence-camembert-large", device="cpu") dataset = datasets.load_dataset("json", data_files=["./dataset.json"], split="train") dataset.load_faiss_index("embeddings", "index.faiss") def search(query, k=3): query_embedding = model.encode(query) _, retrieved_examples = dataset.get_nearest_examples( "embeddings", query_embedding, k=int(k), ) results = [] for text, start, end, title, url in zip( retrieved_examples["text"], retrieved_examples["start"], retrieved_examples["end"], retrieved_examples["title"], retrieved_examples["url"], ): start = start end = end result = { "title": title, "transcript": f"[{str(start)} ====> {str(end)}] {text}", "link": url, } results.append(result) return results iface = gr.Interface( fn=search, inputs=[ gr.inputs.Textbox( label="Query", placeholder="Type in a search query...", lines=3 ), gr.inputs.Number( label="K", default=3, description="Number of results to return", ), ], outputs=[ gr.outputs.Label( label="Result 1", type="auto", default="Search results will appear here." ), gr.outputs.Label( label="Result 2", type="auto", default="" ), gr.outputs.Link( label="Result 3", type="auto", default="" ), ], title="Camembert and Faiss-powered Search Engine", description="Search through a dataset using Camembert and Faiss", theme="default", layout="vertical", allow_flagging=False, allow_screenshot=False, allow_share=True, allow_download=False ) iface.launch()