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

MODEL_NAME = "ibm-granite/granite-20b-code-base-r1.1"

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto"
)


def generate_code(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    outputs = model.generate(
        **inputs,
        max_new_tokens=200,
        temperature=0.7,
        do_sample=True
    )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response


interface = gr.Interface(
    fn=generate_code,
    inputs=gr.Textbox(lines=10, placeholder="Enter your prompt here..."),
    outputs=gr.Textbox(lines=10),
    title="Granite Code Generator",
    description="Generate code using IBM Granite model"
)

if __name__ == "__main__":
    interface.launch()