import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, logging


checkpoint = "Salesforce/codegen-350M-mono"

tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(checkpoint, cache_dir="models/", trust_remote_code=True, revision="main")


def code_gen(text, max_tokens, temp, top_p, rep_penality):
    logging.set_verbosity(logging.CRITICAL)
    pipe = pipeline(
        model=checkpoint,
        max_new_tokens=max_tokens,
        temperature=temp,
        top_p=top_p,
        device= "cuda" if torch.cuda.is_available() else "cpu",
        repetition_penalty=rep_penality
    )

    response = pipe(text)
    print(response)

    return response[0]['generated_text']


Inferece = gr.Interface(
    fn=code_gen,
    inputs=[
        gr.components.Textbox(label="Input what you want, the AI will make the code for you."),
        gr.components.Slider(minimum=128, maximum=512, step=128, label="Choose Max Token"),
        gr.components.Slider(minimum=0.1, maximum=1, step=0.05, label="Choose the model Temperature"),
        gr.components.Slider(minimum=0.1, maximum=1.25, step=0.05, label="Choose top_p"),
        gr.components.Slider(minimum=0.1, maximum=2, step=0.1, label="Choose repetition_penalty")
    ],
    outputs="text",
    live=False   # Ensure live is set to False
)

gr.Markdown('<h2 align="center">AI Code Gen</h2>')
Inferece.launch()