import gradio as gr
from embed import gen_embedding

def generate_embeddings(sequences, plm_model):
    embeddings = gen_embedding(sequences, plm_model=plm_model)
    return embeddings.tolist()

demo = gr.Interface(
    fn=generate_embeddings,
    inputs=[
        "text", 
        gr.Dropdown(choices=["esm1b", "esm2", "prott5", "prostt5"], label="PLM Model")
    ],
    outputs="text"
)
demo.launch()