Spaces:
Sleeping
Sleeping
| """CPU inference. TinyStories keeps backend/local.lua and grammar.lua intact.""" | |
| from functools import lru_cache | |
| import json | |
| import os | |
| from pathlib import Path | |
| ROOT = Path(__file__).resolve().parents[1] | |
| VENDOR = ROOT / 'space' / 'vendor' / 'lua_llama' | |
| TINY_REPO = 'karpathy/tinyllamas' | |
| QWEN_REPO = 'Qwen/Qwen2.5-1.5B-Instruct-GGUF' | |
| def model_file(repo, filename): | |
| from huggingface_hub import hf_hub_download | |
| if repo == TINY_REPO and filename == 'tokenizer.bin': | |
| return str(VENDOR / 'tokenizer.bin') | |
| if repo == TINY_REPO: | |
| local = Path(os.environ.get('LUA_LLAMA_PATH', ROOT.parent / 'lua_llama')) / filename | |
| if local.is_file(): | |
| return str(local) | |
| from space.revisions import REVISIONS | |
| return hf_hub_download(repo, filename, revision=REVISIONS[repo]) | |
| def tiny_weights(device='cpu'): | |
| import torch | |
| from tools.llama2c import load | |
| torch.set_num_threads(min(2, os.cpu_count() or 1)) | |
| cfg, weights = load(model_file(TINY_REPO, 'stories15M.bin')) | |
| return cfg, {name: tensor.to(device) for name, tensor in weights.items()} | |
| def tiny_backend(lua, registry, config): | |
| import torch | |
| from tools.llama2c import forward | |
| cfg, weights = tiny_weights(config.get('device', 'cpu')) | |
| kcache = torch.zeros(cfg.n_layers, cfg.seq_len, cfg.kv_dim, device=weights['tok_emb'].device) | |
| vcache = torch.zeros_like(kcache) | |
| def step(_model, token, pos, *_unused): | |
| return lua.table_from(forward(cfg, weights, token, pos, kcache, vcache).tolist()) | |
| lua.globals().lab_forward = step | |
| lua.globals().lab_cfg = lua.table_from(vars(cfg)) | |
| # Inject only numerical kernels; the original local backend still owns | |
| # tokenization, cache positions, scaffolding, sampling and grammar masks. | |
| lua.execute(''' | |
| package.loaded.model = {new=function() | |
| return {config=lab_cfg, create_run_state=function() return {} end} | |
| end} | |
| package.loaded.generate = {forward=function(...) return lab_forward(...) end} | |
| ''') | |
| options = lua.table_from({ | |
| 'llama_path': VENDOR.as_posix(), | |
| 'tokenizer': model_file(TINY_REPO, 'tokenizer.bin'), | |
| 'tools': registry, 'constrain': config['grammar'], 'temperature': 0.0, | |
| 'max_thought': 24, 'max_call': 96, | |
| }) | |
| backend = lua.eval("require('backend.local').new")(options) | |
| return backend, backend.tok | |
| def qwen_model(): | |
| from llama_cpp import Llama | |
| return Llama( | |
| model_path=model_file(QWEN_REPO, 'qwen2.5-1.5b-instruct-q4_k_m.gguf'), | |
| n_ctx=4096, n_threads=min(2, os.cpu_count() or 1), n_batch=256, | |
| verbose=False, seed=0, | |
| ) | |
| def qwen_backend(lua, registry, config): | |
| model = qwen_model() | |
| encode_json = lua.eval("require('trace').encode") | |
| catalogue = json.loads(encode_json(lua.eval('function(reg) local out={} for _, t in ipairs(reg:list()) do out[#out+1]={name=t.name,args=t.args} end return out end')(registry))) | |
| schemas = [] | |
| for tool in catalogue: | |
| props = {arg['name']: {'type': 'string'} for arg in tool['args']} | |
| required = [arg['name'] for arg in tool['args'] if arg.get('required')] | |
| schemas.append({'type': 'object', 'properties': { | |
| 'tool': {'const': tool['name']}, | |
| 'args': {'type': 'object', 'properties': props, 'required': required, 'additionalProperties': False}, | |
| }, 'required': ['tool', 'args'], 'additionalProperties': False}) | |
| schema = {'oneOf': schemas} | |
| def complete(_self, _prompt, opts): | |
| pieces = json.loads(encode_json(lua.table_from({'system': opts.system, 'steps': opts.steps, 'note': opts.note}))) | |
| messages = [{'role': 'system', 'content': | |
| 'Use the tools to solve the task. Return exactly one JSON object with tool and args. ' | |
| 'All argument values must be strings. Use answer only after observing the required tool results.'}, | |
| {'role': 'user', 'content': pieces['system']}] | |
| for record in pieces.get('steps') or []: | |
| messages.append({'role': 'assistant', 'content': json.dumps(record['call']) if record.get('call') else record.get('raw', '')}) | |
| if record.get('observation'): | |
| messages.append({'role': 'user', 'content': 'RESULT: ' + record['observation']}) | |
| if pieces.get('note'): | |
| messages.append({'role': 'user', 'content': pieces['note']}) | |
| kwargs = {'response_format': {'type': 'json_object', 'schema': schema}} if config['grammar'] else {} | |
| result = model.create_chat_completion(messages=messages, temperature=0, max_tokens=192, **kwargs) | |
| raw = result['choices'][0]['message']['content'] or '' | |
| # Both arms receive identical scaffolding. Free outputs are never repaired. | |
| return 'CALL: ' + raw.replace('\r', ''), lua.table_from({ | |
| 'tokens': result['usage']['completion_tokens'], 'constrained': config['grammar'], | |
| 'finish_reason': result['choices'][0]['finish_reason'], | |
| }) | |
| def encode(_self, text): | |
| return lua.table_from(model.tokenize(text.encode('utf-8'), add_bos=False)) | |
| return lua.table_from({'complete': complete}), lua.table_from({'encode': encode}) | |