import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("SattwikAyyagari/Qwen2.5-Coder-1.5B-NL-Java")
model = AutoModelForCausalLM.from_pretrained("SattwikAyyagari/Qwen2.5-Coder-1.5B-NL-Java")


def generate_code(nl_input: str):
    inputs=tokenizer(f"### Instruction:\n\n{nl_input}\n\n### Response:\n\n",return_tensors="pt")
    outputs=model.generate(**inputs,max_new_tokens=128,eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.eos_token_id)
    decoded=tokenizer.decode(outputs[0],skip_special_tokens=True)
    return decoded.split("### Response:\n\n")[-1].strip()

demo=gr.Interface(
    fn=generate_code,
    inputs=gr.Textbox(label="NL Description"),
    outputs=gr.Textbox(label="Generated JAVA Code"),
    title="CodeGen NL to Java"
)

demo.launch()