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

model_name = "Percy3822/python_ai_coder"  # Replace with your model repo name after training
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")

def generate_code(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=128,
            do_sample=True,
            temperature=0.8,
            top_p=0.95,
            eos_token_id=tokenizer.eos_token_id
        )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

iface = gr.Interface(
    fn=generate_code,
    inputs=gr.Textbox(lines=5, placeholder="Ask me to write/fix/explain Python code..."),
    outputs="text",
    title="Python AI Assistant (Trained on StarCoder)",
    description="Ask it to write functions, fix bugs, explain code, etc."
)

iface.launch()