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

REPO = os.environ.get("ORION_REPO_ID", "Project-Prism/Orion-Flagship")

tokenizer = AutoTokenizer.from_pretrained(REPO, trust_remote_code=True)
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

kwargs = {
    "trust_remote_code": True,
    "dtype": dtype,
}
if torch.cuda.is_available():
    kwargs["device_map"] = "auto"

model = AutoModelForCausalLM.from_pretrained(REPO, **kwargs)
model.eval()


def generate(prompt, max_new_tokens, temperature, top_p):
    inputs = tokenizer(prompt, return_tensors="pt")
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
    with torch.inference_mode():
        ids = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=float(temperature) > 0,
            temperature=max(float(temperature), 1e-5),
            top_p=float(top_p),
            use_cache=False,
        )
    return tokenizer.decode(ids[0], skip_special_tokens=True)


with gr.Blocks(title="Orion Flagship — Project Prism") as demo:
    gr.Markdown(
        "# 🔺 Orion Flagship\n"
        "Experimental **Project Prism / Orion T2** pretraining checkpoint. "
        "This is not an instruction-tuned assistant."
    )
    prompt = gr.Textbox(
        value="The most important reason the sky appears blue is",
        label="Prompt",
        lines=5,
    )
    with gr.Row():
        max_tokens = gr.Slider(1, 128, value=48, step=1, label="New tokens")
        temperature = gr.Slider(0, 1.5, value=0.8, step=0.05, label="Temperature")
        top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.01, label="Top-p")
    go = gr.Button("Generate")
    output = gr.Textbox(label="Output", lines=10)
    go.click(generate, [prompt, max_tokens, temperature, top_p], output)

demo.launch()