Spaces:
Runtime error
Runtime error
File size: 921 Bytes
db24ae6 efd1702 db24ae6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
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() |