import gradio as gr
from transformers import pipeline

pipeline = pipeline("feature-extraction", model="WhereIsAI/UAE-Large-V1")


def predict(text):
    text = text.lower()
    text = text.translate(str.maketrans('', '', string.punctuation))

    
    title_outputs = pipeline(text)
    title_outputs = torch.tensor(title_outputs)

    # Mean pooling
    title_embedding = title_outputs.mean(dim=1)
    #title_embedding = title_outputs.last_hidden_state.mean(dim=1)
    term_outputs = pipeline('multivitamin for men')
    term_outputs = torch.tensor(term_outputs)
    term_embedding = term_outputs.mean(dim=1)
    #term_embedding = term_outputs.last_hidden_state.mean(dim=1)
    semantic_score = cosine_similarity(title_embedding.flatten(), term_embedding.flatten()) * 100
    return title_embedding #str(format(semantic_score,'.2f'))



    
def predict1(text):
    return(text)
    
gradio_app = gr.Interface(
    predict,
    inputs='text',
    outputs='text',
    title="Keyword Score",
    
)

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