import gradio as gr
from transformers import pipeline
from OMEresponse import OMEResponse

pipeline = pipeline(task="text-classification", model="databoyface/distilbert-base-cased-ome-v4.2")

def predict(input: str):
    predictions = pipeline(input)
    result = []
    for p in predictions:
        try:
            response = OMEResponse(p["label"])
        except ValueError:
            response = OMEResponse.INDETERMINATE

        result.append(f"So, possibly {response.reaction} from the guess: \"{response.display_name}\" with score: {p['score']:.6f}.")

    return "\n".join(result)
    
gradio_app = gr.Interface(
    fn=predict,
    inputs=["text"],
    outputs=["text"],
    title="How do you feel?",
)

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