app.py

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()