File size: 4,416 Bytes
29b7b76 24ffcb3 29b7b76 24ffcb3 29b7b76 24ffcb3 29b7b76 36462df 29b7b76 36462df 24ffcb3 36462df 24ffcb3 29b7b76 24ffcb3 29b7b76 28105f9 29b7b76 28105f9 24ffcb3 29b7b76 24ffcb3 29b7b76 24ffcb3 | 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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 | import re
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
MODEL_ID = "UniversalComputingResearch/Atom2.7m"
tokenizer = AutoTokenizer.from_pretrained(
MODEL_ID,
trust_remote_code=True,
)
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
trust_remote_code=True,
).eval()
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
print("Tokenizer class:", type(tokenizer))
print("Model class:", type(model))
print("Device:", device)
def decode_output(output_ids):
"""
First try the official tokenizer.decode path.
If the Space displays raw internal byte-level/LSD tokens, repair display only.
"""
text = tokenizer.decode(output_ids, skip_special_tokens=True)
# If normal decode worked, return it.
if "Ġ" not in text and not re.search(r"\d\s+\d", text):
return text
# Fallback renderer for Atom's internal digit-token representation.
# This does not change model behavior; it only fixes display if the Space
# is not applying the custom decode path correctly.
ids = output_ids.tolist() if hasattr(output_ids, "tolist") else list(output_ids)
ids = [int(x) for x in ids]
digit_id_to_char = {}
for digit in "0123456789":
token_id = tokenizer.convert_tokens_to_ids(digit)
if token_id is not None and token_id != tokenizer.unk_token_id:
digit_id_to_char[int(token_id)] = digit
special_ids = set(tokenizer.all_special_ids)
pieces = []
text_buffer = []
digit_buffer = []
def flush_text():
nonlocal text_buffer
if text_buffer:
part = tokenizer.backend_tokenizer.decode(
text_buffer,
skip_special_tokens=True,
)
part = part.replace("Ġ", " ")
pieces.append(part)
text_buffer = []
def flush_digits():
nonlocal digit_buffer
if digit_buffer:
pieces.extend(reversed(digit_buffer))
digit_buffer = []
for token_id in ids:
if token_id in special_ids:
continue
if token_id in digit_id_to_char:
flush_text()
digit_buffer.append(digit_id_to_char[token_id])
else:
flush_digits()
text_buffer.append(token_id)
flush_text()
flush_digits()
return "".join(pieces).strip()
def generate(prompt, max_new_tokens, temperature, do_sample):
if not prompt.strip():
return "Enter a prompt first."
inputs = tokenizer(
prompt,
return_tensors="pt",
add_special_tokens=False,
)
inputs = {k: v.to(device) for k, v in inputs.items()}
generation_kwargs = {
**inputs,
"max_new_tokens": int(max_new_tokens),
"do_sample": bool(do_sample),
}
if do_sample:
generation_kwargs["temperature"] = float(temperature)
with torch.no_grad():
output_ids = model.generate(**generation_kwargs)
return decode_output(output_ids[0])
examples = [
["12 + 34 =", 3, 1.0, False],
["7 + 8 =", 3, 1.0, False],
["25 - 9 =", 3, 1.0, False],
["5 * 11 =", 3, 1.0, False],
["The capital of France is", 12, 0.8, True],
]
description = """
Atom2.7m is a tiny causal language model for text continuation, with arithmetic-aware handling for numeric spans.
It is not an instruction-tuned chatbot. It works best with short continuation prompts such as `12 + 34 =`.
For arithmetic, greedy decoding with a small number of new tokens usually works best.
"""
demo = gr.Interface(
fn=generate,
inputs=[
gr.Textbox(
label="Prompt",
value="12 + 34 =",
lines=3,
),
gr.Slider(
minimum=1,
maximum=32,
value=3,
step=1,
label="Max new tokens",
),
gr.Slider(
minimum=0.1,
maximum=2.0,
value=1.0,
step=0.1,
label="Temperature",
),
gr.Checkbox(
value=False,
label="Sample instead of greedy decoding",
),
],
outputs=gr.Textbox(
label="Model output",
lines=6,
),
title="Atom2.7m Arithmetic Demo",
description=description,
examples=examples,
allow_flagging="never",
)
if __name__ == "__main__":
demo.launch() |