File size: 1,837 Bytes
6f2679b eca12d4 6f2679b |
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 |
import os, torch, gradio as gr
from threading import Thread
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1") # faster download on Spaces
MODEL_ID = "TildeAI/TildeOpen-30b"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False)
# load in BF16 and let HF map devices automatically
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
torch_dtype=torch.bfloat16,
device_map="auto",
)
# slight speedups on A100
torch.backends.cuda.matmul.allow_tf32 = True
SYS = (
"You are a helpful multilingual assistant. "
"This is a *base* model (not instruction tuned); follow the user's request precisely."
)
def build_prompt(history, user_msg):
# simple conversation transcript; base models don't need a special chat template
parts = [SYS, ""]
for u, a in history:
parts += [f"User: {u}", f"Assistant: {a}"]
parts += [f"User: {user_msg}", "Assistant:"]
return "\n".join(parts)
def chat_fn(message, history):
prompt = build_prompt(history, message)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
gen_kwargs = dict(
**inputs,
max_new_tokens=512,
do_sample=True,
temperature=0.7,
top_p=0.9,
repetition_penalty=1.1,
streamer=streamer,
)
t = Thread(target=model.generate, kwargs=gen_kwargs)
t.start()
partial = ""
for chunk in streamer:
partial += chunk
yield partial
demo = gr.ChatInterface(
fn=chat_fn,
title="TildeOpen-30B (Transformers, BF16)",
description="Base model; multilingual. If build fails with OOM, switch to Option B (GGUF).",
)
demo.queue().launch()
|