File size: 1,884 Bytes
0975e0c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | 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()
|