import gradio as gr
import torch
from transformers import AutoTokenizer, GenerationConfig, GPT2Config, GPT2LMHeadModel

# 1. Load the standard GPT-2 tokenizer
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

# 3. Load the model 
config = GPT2Config.from_pretrained(
    pretrained_model_name_or_path = "gpt2",
    vocab_size = len(tokenizer),
    n_ctx = 256,
    bos_token_id = tokenizer.bos_token_id,
    eos_token_id = tokenizer.bos_token_id
)
model = GPT2LMHeadModel(config)

def generate_code(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            pad_token_id=tokenizer.eos_token_id
        )
    
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# 5. Gradio Interface
demo = gr.Interface(
    fn=generate_code,
    inputs=gr.Textbox(placeholder="Write a function to...", label="Input Prompt"),
    outputs=gr.Code(label="GPT-2 Generated Code", language="python"),
    title="Small Code Snippet Generator"
)

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