import gradio as gr
from utils_model import ModelFactory

class GradioAppManager():

    def __init__(self, model_type):
        model_factory_obj = ModelFactory()
        self.model = model_factory_obj.create_model(model_type)

    def __retrieve_embeddings__(self, input_queries_df, limitnoffeatures):
        queries_list = input_queries_df.values
        queries_embeddings_list = self.model.retrieve_embeddings_from_texts_list(queries_list, limitnoffeatures=limitnoffeatures)

        return queries_embeddings_list

    def build(self):
        gr_input_dataframe = gr.Dataframe(headers=['queries'], datatype=['str'], row_count=2, col_count=(1, 'fixed'))
        gr_number_limitnoffeatures = gr.Number(50)

        app = gr.Interface(fn=self.__retrieve_embeddings__, 
                        inputs=[gr_input_dataframe, gr_number_limitnoffeatures], 
                        outputs="dataframe")

        return app

    def run(self, app):
        app.launch(server_name='0.0.0.0')

gradio_app_manager_obj = GradioAppManager('multilingual-e5-large-ft-sts-spanish-matryoshka-768-64-5e')
app = gradio_app_manager_obj.build()
gradio_app_manager_obj.run(app)