import gradio as gr
from transformers import pipeline

pipe = pipeline(task="text2text-generation", model="Jumpy-pku/t5-dish-name-recognition")
pipe_aux = pipeline(task="text2text-generation", model="Jumpy-pku/t5-dish-name-recognition-auxiliary")

def predict(input_text):
    
    verbs = pipe_aux(f"“{input_text}”这个菜谱的主要操作是什么？", max_length=20, num_return_sequences=1)[0]["generated_text"]
    flavs = pipe_aux(f"“{input_text}”这个菜谱的主要风味是什么？", max_length=20, num_return_sequences=1)[0]["generated_text"]
    ings = pipe_aux(f"“{input_text}”这个菜谱的主要食材是什么？", max_length=20, num_return_sequences=1)[0]["generated_text"]

    comp_prompt = ""
    if verbs:
        comp_prompt += f"这个菜谱的主要操作是{verbs},"
    if flavs:
        comp_prompt += f"这个菜谱的主要风味是{flavs},"
    if ings:
        comp_prompt += f"这个菜谱的主要食材是{ings},"
    comp_prompt += "这个菜谱的菜名是什么？"

    return pipe(f"“{input_text}”{comp_prompt}", max_length=20, num_return_sequences=1)[0]["generated_text"]

    
demo = gr.Interface(fn=predict, inputs="text", outputs="text", title="菜名生成", description="Input the instructions of a Chinese recipe, the model will generate the corresponding dish name.\n输入菜谱步骤，生成菜名。")
demo.launch()