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

MODEL_NAME = "tiiuae/Falcon3-7B-Base"

access_token = os.getenv("HF_ACCESS_TOKEN")
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME, 
    trust_remote_code=True,
    token=access_token)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
    token=access_token
)

def improve_code(code: str) -> str:
    prompt = (
        "You are an expert code assistant.\n"
        "Given the following code, suggest an improved version with clear comments and best practices.\n"
        "Output only the improved code.\n\n"
        f"Original code:\n{code}\n\nImproved code:"
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            input_ids=inputs["input_ids"],
            attention_mask=inputs["attention_mask"],
            max_new_tokens=512,
            temperature=0.2,
            top_p=0.9,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
    generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
    improved = generated.split("Improved code:")[-1].strip()
    return improved

app = gr.Blocks()

with app:
    gr.Markdown("## MCP Server Code Improver with Falcon3-7B-Base")
    code_input = gr.Textbox(label="Original code", lines=15)
    improve_btn = gr.Button("Improve Code")
    code_output = gr.Textbox(label="Improved code", lines=15)
    improve_btn.click(improve_code, inputs=code_input, outputs=code_output)

app.launch()