search / app.py
Ubuntu
v1
efd1702
raw
history blame contribute delete
921 Bytes
from sentence_transformers import SentenceTransformer
import hnswlib
import pandas as pd
import gradio as gr
model = SentenceTransformer('rufimelo/Legal-SBERTimbau-sts-base')
videos_db = pd.read_csv('videos_db.tsv', header=None, names=["id", "title", "thumb"], sep='\t')
video_index = hnswlib.Index('cosine', dim=model.get_sentence_embedding_dimension())
video_index.load_index('index.bin')
def predict(query):
query_embedding = model.encode([query])
i, _ = video_index.knn_query(query_embedding, k=5)
code = "<table><tbody>"
for e in videos_db.iloc[i[0]].values:
code += f'<tr><td><img height="90" width="150" src="{e[2]}"></td><td><a style="font-family:Arial" href="https://www.youtube.com/watch?v={e[0]}">{e[1]}</a></td></tr>'
code += "</tbody></table>"
return code
demo = gr.Interface(fn=predict,
inputs="text",
outputs="html")
demo.launch()