import gradio as gr
#from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers import pipeline

model_name = "vijjuk/codegen-350M-mono-python-18k-alpaca"
pipe = pipeline("python-fine-tuning", model=model_name)





#base_model = AutoModelForCausalLM.from_pretrained(model_name)
#tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
#tokenizer.pad_token = tokenizer.eos_token
#tokenizer.padding_side = "right"

def query(instruction, input):
    prompt = f"""### Instruction:
        Use the Task below and the Input given to write the Response, which is a programming code that can solve the Task.
        ### Task:
        {instruction}
        ### Input:
        {input}
        ### Response:
        """
    #input_ids = tokenizer(prompt, return_tensors="pt", truncation=True)
    #output_base = base_model.generate(input_ids=input_ids, max_new_tokens=500, do_sample=True, top_p=0.9,temperature=0.5)
    #response = "{tokenizer.batch_decode(output_base.detach().cpu().numpy(), skip_special_tokens=True)[0][len(prompt):]}"
    #return response
    return pipe(prompt)[0]["prompt"]

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