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="Enter your request, and the AI will generate the code for you."),
        gr.components.Slider(minimum=128, maximum=1024, step=128, value=512, label="Choose Max Token Size"),
        gr.components.Slider(minimum=0.1, maximum=1, step=0.05, value=0.65, label="Choose the model Temperature"),
        gr.components.Slider(minimum=0.1, maximum=1.25, step=0.05, value=0.9, label="Choose top_p"),
        gr.components.Slider(minimum=0.1, maximum=2, step=0.1, value=1.15, label="Choose repetition_penalty")
    ],
    outputs="text",
    title="AI Code Gen",
    live=False
)

Inferece.queue(concurrency_count=1)
Inferece.launch()