import gradio as gr
from transformers import pipeline
title= "German Flan-T5"
examples = [
  ["Erzähl mit eine Geschichte!",200,2,3,10,"Deutsch"]
]

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=200.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=20.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"]),
  ],
  outputs="text",
  title=title,
  examples=examples
)

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