elise/src/app.py

"""
Gradio app builder
"""
import gradio as gr
from sentence_transformers import SentenceTransformer

from nlu import SentenceParser
from matcher import Matcher
from utils import df_to_json


# prep models
MODEL_CHECKPOINT = (
    "tner/roberta-large-mit-restaurant"  # "BerserkerMother/restaurant_ner"
)
parser = SentenceParser.from_huggingface(MODEL_CHECKPOINT)

embedder = SentenceTransformer("all-MiniLM-L6-v2")
matcher = Matcher.from_path("data/final_data.csv", embedder)


def recommend(query: str):
    """
    App logic for gradio fn

        Parameters:
            query(str): user query
    """
    ner_tags = parser.get_ner(query)
    jobs = parser.get_jobs(ner_tags)
    recomms = matcher.handle_jobs(jobs)
    recomms = df_to_json(recomms, [query])
    return {"jobs": jobs, "recomms": recomms}


iface = gr.Interface(
    fn=recommend,
    inputs="text",
    outputs="json",
    theme=gr.themes.Default(
        primary_hue=gr.themes.colors.rose,
        neutral_hue=gr.themes.colors.gray,
    ),
)
iface.launch()