import gradio as gr
from peft import get_peft_config, get_peft_model, PeftModel, PeftConfig, LoraConfig, TaskType

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# huggingface hub model id
model_id="xavierbarbier/flan-t5-small-ameli_qa_10k"


# load model from the hub
model = AutoModelForSeq2SeqLM.from_pretrained(model_id)

tokenizer = AutoTokenizer.from_pretrained(model_id)

max_target_length = 512

def greet(input_text):
  input_ids = tokenizer(input_text, return_tensors="pt").input_ids
  outputs = model.generate(input_ids = input_ids, max_length = max_target_length)
  answer = tokenizer.decode(outputs[0]).replace("<pad> ","").replace("</s>","")

  return answer

iface = gr.Interface(fn=greet, inputs=["text"],
                     outputs="text")
iface.launch()