File size: 1,369 Bytes
45d6b11
2014880
 
 
 
45d6b11
2014880
45d6b11
 
 
 
a6fecbc
45d6b11
2014880
45d6b11
 
2014880
 
 
 
45d6b11
2014880
45d6b11
2014880
 
 
 
 
 
 
 
 
45d6b11
 
 
 
 
 
 
2014880
 
45d6b11
 
 
 
 
 
 
 
 
2014880
 
45d6b11
 
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
import logging
import gradio as gr
import datasets
import sentence_transformers

logging.disable(logging.CRITICAL)

model = sentence_transformers.SentenceTransformer(
    "dangvantuan/sentence-camembert-large", device="cuda"
)

dataset = datasets.load_dataset("json", data_files=["./dataset.json"], split="train")
dataset.load_faiss_index("embeddings", "index.faiss")

def search(query: str, k: int):
    query_embedding = model.encode(query)
    _, retrieved_examples = dataset.get_nearest_examples(
        "embeddings",
        query_embedding,
        k=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=["text", "number"],
    outputs=gr.outputs.JSON(),
    title="Search Dataset",
    description="Search a dataset using Camembert and Faiss.",
    example=[
        "Enter a query to search for.",
        5
    ]
)

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