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()