import gradio as gr
from transformers import AutoModel, AutoConfig
from main_idea_with_torch import predict_mainidea_sent_old
from main_idea_with_pipeline import predict_mainidea_sent

config = AutoConfig.from_pretrained("yutingg/custom-distill-bert-for-sentence-label", trust_remote_code=True)
model = AutoModel.from_pretrained("yutingg/custom-distill-bert-for-sentence-label", trust_remote_code=True, config=config)

def greet(essay):
    ret = predict_mainidea_sent(essay, model), predict_mainidea_sent_old(essay, model)
    return ret

iface = gr.Interface(fn=greet, inputs="text", outputs=[
    gr.Dataframe(
        label="pipeline output",
        headers=['label: is main idea', 'sentence'],
        datatype=["str", "str"],
        col_count=(2, "fixed"),
    ),
    gr.Dataframe(
        label="torch output with Triage",
        headers=['label: is main idea', 'sentence'],
        datatype=["str", "str"],
        col_count=(2, "fixed"),
    )
])
iface.launch()