Spaces:
Runtime error
Runtime error
File size: 3,047 Bytes
880bf35 7d7ef5b 880bf35 |
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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
import os
os.environ.setdefault("HF_HOME", "/tmp/hf")
os.environ.setdefault("HF_HUB_CACHE", "/tmp/hf/hub")
os.environ.setdefault("TRANSFORMERS_CACHE", "/tmp/hf/transformers")
os.environ.setdefault("NANOCHAT_BASE_DIR", "/tmp/nanochat")
from huggingface_hub import hf_hub_download
import torch
import gradio as gr
from nanochat.checkpoint_manager import load_model_from_dir
from nanochat.engine import Engine
# Hardcoded model selection for this Space
MODEL_REPO = "loocorez/nanochat-base-d20-step21400"
STEP = "021400"
DEPTH = "20"
ckpt_dir = f"/tmp/ckpt/d{DEPTH}"
os.makedirs(ckpt_dir, exist_ok=True)
# tokenizer (where nanochat expects it)
tokenizer_dir = "/tmp/nanochat/tokenizer"
os.makedirs(tokenizer_dir, exist_ok=True)
hf_hub_download(MODEL_REPO, "tokenizer/tokenizer.pkl", local_dir=tokenizer_dir, local_dir_use_symlinks=False)
# base checkpoint
hf_hub_download(MODEL_REPO, f"base_checkpoints/d{DEPTH}/model_{STEP}.pt", local_dir=ckpt_dir, local_dir_use_symlinks=False)
hf_hub_download(MODEL_REPO, f"base_checkpoints/d{DEPTH}/meta_{STEP}.json", local_dir=ckpt_dir, local_dir_use_symlinks=False)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model, tokenizer, _ = load_model_from_dir(ckpt_dir, device, phase="eval")
engine = Engine(model, tokenizer)
def chat_fn(history, temperature=0.8, top_k=50, max_new_tokens=256):
bos = tokenizer.get_bos_token_id()
user_start = tokenizer.encode_special("<|user_start|>")
user_end = tokenizer.encode_special("<|user_end|>")
assistant_start = tokenizer.encode_special("<|assistant_start|>")
assistant_end = tokenizer.encode_special("<|assistant_end|>")
tokens = [bos]
for role, content in history:
if role == "user":
tokens += [user_start] + tokenizer.encode(content) + [user_end]
else:
tokens += [assistant_start] + tokenizer.encode(content) + [assistant_end]
tokens += [assistant_start]
with torch.amp.autocast(device_type="cuda" if device.type == "cuda" else "cpu", dtype=torch.bfloat16 if device.type == "cuda" else torch.float32):
token_column, _ = next(engine.generate(tokens, num_samples=1, max_tokens=max_new_tokens, temperature=temperature, top_k=top_k))
new_tokens = token_column[len(tokens):]
return tokenizer.decode(new_tokens)
with gr.Blocks() as demo:
gr.Markdown("# NanoChat BASE")
chat = gr.Chatbot(type="tuple")
msg = gr.Textbox()
temp = gr.Slider(0.0, 1.5, value=0.8, step=0.05, label="Temperature")
topk = gr.Slider(1, 200, value=50, step=1, label="Top-k")
max_toks = gr.Slider(16, 1024, value=256, step=16, label="Max new tokens")
def respond(user_message, chat_history, temperature, top_k, max_new_tokens):
chat_history = chat_history + [("user", user_message)]
reply = chat_fn(chat_history, temperature, top_k, max_new_tokens)
chat_history = chat_history + [("assistant", reply)]
return "", chat_history
msg.submit(respond, [msg, chat, temp, topk, max_toks], [msg, chat])
demo.launch()
|