import gradio as gr
from transformers import pipeline
title= "German Flan-T5"
desc="Kommunikation mit flan-t5-large auf Deutsch wird intern ins Englische (opus-mt-de-en) und vom Englischen (opus-mt-en-de) übersetzt."
examples = [
  ["Erzähl mit eine Geschichte!",50,2,3,1,"Deutsch"],
  ["Welche Blumen sollte man jemandem zum Valentinstag schenken?",50,1,0,1,"Deutsch"],  
  ["Please write a step by step recipe to make bolognese pasta!",50,2,3,2,"Englisch"]
]

tDeEn = pipeline(model="Helsinki-NLP/opus-mt-de-en")
tEnDe = pipeline(model="Helsinki-NLP/opus-mt-en-de")
bot = pipeline(model="google/flan-t5-large")

def solve(text,max_length,length_penalty,no_repeat_ngram_size,num_beams,language):
  if(language=="Deutsch"): 
      text=tDeEn(text)[0]["translation_text"]
  out=bot(text,max_length=max_length, length_penalty=length_penalty, no_repeat_ngram_size=no_repeat_ngram_size, num_beams=num_beams, early_stopping=True)[0]["generated_text"]
  if(language=="Deutsch"): 
      out=tEnDe(out)[0]["translation_text"]
  return out

task = gr.Interface(
  fn=solve,
  inputs=[
      gr.Textbox(lines=5,max_lines=6,label="Frage"),
      gr.Slider(minimum=1.0,maximum=200.0,value=50.0,step=1,interactive=True,label="max_length"),
      gr.Slider(minimum=1.0,maximum=20.0,value=1.0,step=1,interactive=True,label="length_penalty"),
      gr.Slider(minimum=0.0,maximum=5.0,value=3.0,step=1,interactive=True,label="no_repeat_ngram_size"),
      gr.Slider(minimum=1.0,maximum=20.0,value=1.0,step=1,interactive=True,label="num_beams"),
      gr.Dropdown(["Deutsch", "Englisch"],value="Deutsch"),
  ],
  outputs="text",
  title=title,
  description=desc,
  examples=examples
)

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