diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..8b1e4f74717a0931f3c21a0da546c03d2af94bea 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,47 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +assert/ocr.jpg filter=lfs diff=lfs merge=lfs -text +assert/person.jpg filter=lfs diff=lfs merge=lfs -text +assert/phrase_grounding.jpg filter=lfs diff=lfs merge=lfs -text +image_encoder_mlp.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2.5_tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l0_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l10_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l11_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l12_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l13_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l14_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l15_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l16_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l17_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l18_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l19_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l1_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l20_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l21_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l22_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l23_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l24_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l25_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l26_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l27_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l28_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l29_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l2_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l30_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l31_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l32_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l33_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l34_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l35_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l3_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l4_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l5_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l6_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l7_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l8_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_p128_l9_together.axmodel filter=lfs diff=lfs merge=lfs -text +qwen2_post.axmodel filter=lfs diff=lfs merge=lfs -text +test_data/ocr.jpg filter=lfs diff=lfs merge=lfs -text +test_data/person.jpg filter=lfs diff=lfs merge=lfs -text diff --git a/README.md b/README.md index 555b4aed7fba6658868d69c9844d1dcb224a0d31..1a828371511f5aaa3b62bf5146f112ff1ce82dad 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,167 @@ --- -license: bsd-3-clause +license: mit +language: +- en +- zh +base_model: +- nvidia/LocateAnything-3B +pipeline_tag: detection, VLM +library_name: transformers +tags: +- LocateAnything-3B +- Int4 +- VLM +- GPTQ --- + +# LocateAnything-3B + +This version of LocateAnything-3B have been converted to run on the Axera NPU using **w4a16** quantization. + +Compatible with Pulsar2 version: 6.0 + +## Convert tools links: + +For those who are interested in model conversion, you can try to export axmodel through the original repo : + +- https://huggingface.co/nvidia/LocateAnything-3B + +[Pulsar2 Link, How to Convert LLM from Huggingface to axmodel](https://pulsar2-docs.readthedocs.io/en/latest/appendix/build_llm.html) + +[AXera NPU HOST LLM Runtime](TODO) + + +## Support Platform + +- AX650 + - AX650N DEMO Board + - [M4N-Dock(爱芯派Pro)](https://wiki.sipeed.com/hardware/zh/maixIV/m4ndock/m4ndock.html) + - [M.2 Accelerator card](https://docs.m5stack.com/zh_CN/ai_hardware/LLM-8850_Card) + +**Image Process** +|Chips| input size | image num | image encoder | ttft(493 tokens) | w4a16 | CMM | Flash | +|--|--|--|--|--|--|--|--| +|AX650| 560*560 | 1 | 1152.583 ms | 2072.06 ms | 10.61 tokens/sec| 2.9GiB | 3.2GiB | + +The DDR capacity refers to the CMM memory that needs to be consumed. Ensure that the CMM memory allocation on the development board is greater than this value. + +## How to use + +## 安装 axllm +方式一：克隆仓库后执行安装脚本： + +```shell +git clone -b axllm https://github.com/AXERA-TECH/ax-llm.git +cd ax-llm +./install.sh +``` + +方式二：一行命令安装（默认分支 `axllm`）： + +```shell +curl -fsSL https://raw.githubusercontent.com/AXERA-TECH/ax-llm/axllm/install.sh | bash +``` + +方式三：下载Github Actions CI 导出的可执行程序（适合没有编译环境的用户）： + +如果没有编译环境，请到： +`https://github.com/AXERA-TECH/ax-llm/actions?query=branch%3Aaxllm` +下载 **最新 CI 导出的可执行程序**（`axllm`），然后： + +```shell +chmod +x axllm +sudo mv axllm /usr/bin/axllm +``` + +## 模型下载（Hugging Face） +先创建模型目录并进入，然后下载到该目录： + +```shell +mkdir -p AXERA-TECH/LocateAnything-3B +cd AXERA-TECH/LocateAnything-3B +hf download AXERA-TECH/LocateAnything-3B --local-dir . + +# structure of the downloaded files +tree -L 3 +. +└── AXERA-TECH + └── LocateAnything-3B + |-- assert + |-- config.json + |-- gradio_locateanything_axengine.py + |-- image_encoder_mlp.axmodel + |-- infer_locateanything_axengine.py + |-- model.embed_tokens.weight.bfloat16.bin + |-- post_config.json + |-- qwen2.5_tokenizer + |-- qwen2_5_tokenizer.txt + |-- qwen2_p128_l0_together.axmodel + |-- qwen2_p128_l10_together.axmodel + |-- qwen2_p128_l11_together.axmodel + |-- qwen2_p128_l12_together.axmodel + |-- qwen2_p128_l13_together.axmodel + |-- qwen2_p128_l14_together.axmodel + |-- qwen2_p128_l15_together.axmodel + |-- qwen2_p128_l16_together.axmodel + |-- qwen2_p128_l17_together.axmodel + |-- qwen2_p128_l18_together.axmodel + |-- qwen2_p128_l19_together.axmodel + |-- qwen2_p128_l1_together.axmodel + |-- qwen2_p128_l20_together.axmodel + |-- qwen2_p128_l21_together.axmodel + |-- qwen2_p128_l22_together.axmodel + |-- qwen2_p128_l23_together.axmodel + |-- qwen2_p128_l24_together.axmodel + |-- qwen2_p128_l25_together.axmodel + |-- qwen2_p128_l26_together.axmodel + |-- qwen2_p128_l27_together.axmodel + |-- qwen2_p128_l28_together.axmodel + |-- qwen2_p128_l29_together.axmodel + |-- qwen2_p128_l2_together.axmodel + |-- qwen2_p128_l30_together.axmodel + |-- qwen2_p128_l31_together.axmodel + |-- qwen2_p128_l32_together.axmodel + |-- qwen2_p128_l33_together.axmodel + |-- qwen2_p128_l34_together.axmodel + |-- qwen2_p128_l35_together.axmodel + |-- qwen2_p128_l3_together.axmodel + |-- qwen2_p128_l4_together.axmodel + |-- qwen2_p128_l5_together.axmodel + |-- qwen2_p128_l6_together.axmodel + |-- qwen2_p128_l7_together.axmodel + |-- qwen2_p128_l8_together.axmodel + |-- qwen2_p128_l9_together.axmodel + |-- qwen2_post.axmodel + |-- results + `-- test_data + +4 directories, 44 files +``` + +## Inference with AX650 Host, such as M4N-Dock(爱芯派Pro) or AX650N DEMO Board + +### Gradio Demo + +```shell +(base) root@ax650:~/LocateAnything# python gradio_locateanything_axengine.py +[INFO] Available providers: ['AxEngineExecutionProvider', 'AXCLRTExecutionProvider'] +[Gradio] starting LocateAnything UI +[Gradio] local: http://127.0.0.1:7860 +[Gradio] LAN: http://10.126.29.50:7860 +[Gradio] LAN: http://10.126.29.68:7860 +[Gradio] LAN: http://172.17.0.1:7860 +[Gradio] Use another computer in the same LAN to open the LAN URL. +* Running on local URL: http://0.0.0.0:7860 +* To create a public link, set `share=True` in `launch()`. +``` + +Output: + +detection: +![detection](./assert/person.jpg) + +ocr: +![ocr](./assert/ocr.jpg) + +phrase grounding: +![phrase grounding](./assert/phrase_grounding.jpg) diff --git a/assert/ocr.jpg b/assert/ocr.jpg new file mode 100644 index 0000000000000000000000000000000000000000..19fd97d33a329642381fa1477d8ed9b39a76125b --- /dev/null +++ b/assert/ocr.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1343d9c0b72463b5b0c86aedc19478f6bc1ff45e541686852e8d641e57303b2c +size 251985 diff --git a/assert/person.jpg b/assert/person.jpg new file mode 100644 index 0000000000000000000000000000000000000000..7325ca84d40b679fb75fb251d96f522d7c2db5a0 --- /dev/null +++ b/assert/person.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9e244a1550a914272e6929d87a814bb0da29a538a77a24b793917b9bd18844c5 +size 677828 diff --git a/assert/phrase_grounding.jpg b/assert/phrase_grounding.jpg new file mode 100644 index 0000000000000000000000000000000000000000..0780a9557a3503a1d00776d43102abd42a97c7e7 --- /dev/null +++ b/assert/phrase_grounding.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f30e3af96d68f93e485ec03661e697ba6a22fc537095233d6ff788d88ccb029f +size 313152 diff --git a/config.json b/config.json new file mode 100644 index 0000000000000000000000000000000000000000..c1368d7624c57fef398a9de4951a025a5589a769 --- /dev/null +++ b/config.json @@ -0,0 +1,16 @@ +{ + "system_prompt": "you are a helpful assistant.", + "model_name": "AXERA-TECH/Qwen2.5-locateanything", + "url_tokenizer_model": "qwen2_5_tokenizer.txt", + "tokenizer_type": "Qwen2_5", + "post_config_path": "post_config.json", + "template_filename_axmodel": "qwen2_p128_l%d_together.axmodel", + "axmodel_num": 36, + "filename_post_axmodel": "qwen2_post.axmodel", + "filename_tokens_embed": "model.embed_tokens.weight.bfloat16.bin", + "tokens_embed_num": 152681, + "tokens_embed_size": 2048, + "use_mmap_load_embed": true, + "use_mmap_load_layer": true, + "devices": [0] +} diff --git a/gradio_locateanything_axengine.py b/gradio_locateanything_axengine.py new file mode 100644 index 0000000000000000000000000000000000000000..76ca6f88bc005cc052966466391b20d5a9696bee --- /dev/null +++ b/gradio_locateanything_axengine.py @@ -0,0 +1,477 @@ +#!/usr/bin/env python3 +"""Gradio UI for LocateAnything axengine inference. + +The UI reuses infer_locateanything_axengine.py and keeps the same pure +axengine path: image encoder + AR LLM decode + strict streaming geometry +decode. Gradio is intentionally imported lazily so the script can print a +clear dependency error on boards where it is not preinstalled. +""" + +from __future__ import annotations + +import argparse +import json +import os +import queue +import re +import subprocess +import socket +import threading +import time +from pathlib import Path +from types import SimpleNamespace + +from PIL import Image, UnidentifiedImageError + +from infer_locateanything_axengine import ( + DEFAULT_IMAGE_ENCODER, + DEFAULT_LLM_DIR, + DEFAULT_SYSTEM_PROMPT, + DEFAULT_TARGET, + DEFAULT_TOKENIZER, + LocateAnythingAxEngineRunner, + GeometryImageDrawer, + PROMPT_SPECS, + Point, + Box, + StreamingGeometryDecoder, + build_task_prompt, + time_geometry_decode, +) + + +SCRIPT_DIR = Path(__file__).resolve().parent +UI_DIR = SCRIPT_DIR / "results" / "gradio" +INPUT_DIR = UI_DIR / "inputs" +OUTPUT_DIR = UI_DIR / "outputs" + +UI_TASKS = { + "Object Detection": "object_detection", + "Phrase Grounding (Single Box)": "phrase_grounding_single", + "Phrase Grounding (Multiple Boxes)": "phrase_grounding_multi", + "Text Grounding": "text_grounding", + "Scene Text Detection (OCR)": "scene_text_detection", + "Document Layout Analysis": "document_layout", + "GUI Grounding (Box)": "gui_grounding_box", + "GUI Grounding (Point)": "gui_grounding_point", + "Pointing": "pointing", +} + +gr = None + + +def ensure_gradio() -> None: + global gr + if gr is not None: + return + try: + import gradio as gradio_module + except ModuleNotFoundError: + print("[ERROR] gradio is not installed.") + print("Install it first, for example: python3 -m pip install gradio") + raise SystemExit(1) + gr = gradio_module + + +class RunnerState: + def __init__(self, tokenizer: str, llm_dir: str, image_encoder: str): + self.tokenizer = tokenizer + self.llm_dir = llm_dir + self.image_encoder = image_encoder + self.runner: LocateAnythingAxEngineRunner | None = None + self.lock = threading.Lock() + + def get_runner(self) -> LocateAnythingAxEngineRunner: + if self.runner is None: + self.runner = LocateAnythingAxEngineRunner(self.tokenizer, self.llm_dir, self.image_encoder) + return self.runner + + def close(self) -> None: + if self.runner is not None: + self.runner.close() + self.runner = None + + +def get_lan_ips() -> list[str]: + ips = set() + try: + output = subprocess.check_output(["hostname", "-I"], text=True, timeout=2) + for item in output.split(): + if item and not item.startswith("127.") and ":" not in item: + ips.add(item) + except (OSError, subprocess.SubprocessError): + pass + + try: + output = subprocess.check_output(["ip", "-4", "addr", "show"], text=True, timeout=2) + for line in output.splitlines(): + line = line.strip() + if not line.startswith("inet "): + continue + item = line.split()[1].split("/")[0] + if item and not item.startswith("127."): + ips.add(item) + except (OSError, subprocess.SubprocessError): + pass + + try: + hostname = socket.gethostname() + for item in socket.gethostbyname_ex(hostname)[2]: + if item and not item.startswith("127."): + ips.add(item) + except OSError: + pass + + try: + sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + sock.connect(("8.8.8.8", 80)) + ip = sock.getsockname()[0] + if ip and not ip.startswith("127."): + ips.add(ip) + sock.close() + except OSError: + pass + + return sorted(ips) or ["127.0.0.1"] + + +def find_available_port(start_port: int, retries: int) -> int: + for port in range(start_port, start_port + retries + 1): + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", port)) + return port + except OSError: + continue + raise OSError(f"Cannot find empty port in range: {start_port}-{start_port + retries}") + + +def save_uploaded_image(image: Image.Image) -> str: + INPUT_DIR.mkdir(parents=True, exist_ok=True) + ts = time.strftime("%Y%m%d_%H%M%S") + path = INPUT_DIR / f"upload_{ts}_{int(time.time() * 1000) % 1000:03d}.jpg" + image.convert("RGB").save(path) + return str(path) + + +def load_ui_image(path: str) -> Image.Image | None: + if not path or not os.path.exists(path): + return None + for _ in range(5): + try: + with Image.open(path) as image: + return image.convert("RGB").copy() + except (OSError, UnidentifiedImageError): + time.sleep(0.05) + return None + + +def drain_update_queue(update_queue: queue.Queue) -> None: + while True: + try: + update_queue.get_nowait() + except queue.Empty: + return + + +def build_prompt(ui_task: str, target_text: str, custom_prompt: str) -> tuple[str, str, str, str]: + task = UI_TASKS[ui_task] + if custom_prompt.strip(): + ns = SimpleNamespace(prompt=custom_prompt.strip(), task=task, target=target_text.strip(), categories=None, phrase=None) + return build_task_prompt(ns) + + target = target_text.strip() or DEFAULT_TARGET + spec = PROMPT_SPECS[task] + categories = target if spec.target_kind == "categories" else None + phrase = target if spec.target_kind == "phrase" else None + ns = SimpleNamespace(prompt=None, task=task, target=target, categories=categories, phrase=phrase) + return build_task_prompt(ns) + + +def box_to_dict(box: Box) -> dict[str, float]: + return {"x1": box.x1, "y1": box.y1, "x2": box.x2, "y2": box.y2} + + +def point_to_dict(point: Point) -> dict[str, float]: + return {"x": point.x, "y": point.y} + + +def format_coord(value: float) -> int: + return int(round(value)) + + +def extract_ref_labels(text: str) -> list[str]: + refs = re.findall(r"(.*?)\s*", text, flags=re.DOTALL) + return [ref.strip() for ref in refs] + + +def format_detection_results( + label: str, + boxes: list[Box], + points: list[Point], + ref_labels: list[str] | None = None, +) -> str: + fallback_ref = label or "text" + lines = [] + ref_labels = ref_labels or [] + for index, box in enumerate(boxes): + ref = ref_labels[index] if index < len(ref_labels) and ref_labels[index] else fallback_ref + box_value = [format_coord(box.x1), format_coord(box.y1), format_coord(box.x2), format_coord(box.y2)] + lines.append(json.dumps([ref, {"box": box_value}], ensure_ascii=False)) + for point in points: + point_value = [format_coord(point.x), format_coord(point.y)] + lines.append(json.dumps([fallback_ref, {"point": point_value}], ensure_ascii=False)) + return "\n".join(lines) if lines else "[]" + + +def make_json_result( + *, + task: str, + output_type: str, + target: str, + prompt: str, + token_ids: list[int], + text: str, + boxes: list[Box], + points: list[Point], + stream_boxes: list[dict], + stream_points: list[dict], + stream_consistent: bool, + stream_points_consistent: bool, + timings: dict, + output_image: str, +) -> dict: + return { + "task": task, + "task_output_type": output_type, + "target": target, + "prompt": prompt, + "token_ids": token_ids, + "text": text, + "boxes": [box_to_dict(box) for box in boxes], + "points": [point_to_dict(point) for point in points], + "stream_boxes": stream_boxes, + "stream_points": stream_points, + "stream_consistent": stream_consistent, + "stream_points_consistent": stream_points_consistent, + "timings": timings, + "output_image": output_image, + } + + +def create_infer_fn(state: RunnerState, args: argparse.Namespace): + def infer( + image: Image.Image | None, + ui_task: str, + target_text: str, + custom_prompt: str, + max_new_tokens: int, + temperature: float, + top_p: float, + repetition_penalty: float, + seed: int, + ): + if image is None: + yield "请先上传图片。", None, None + return + + OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + image_path = save_uploaded_image(image) + ts = time.strftime("%Y%m%d_%H%M%S") + stem = f"{UI_TASKS[ui_task]}_{ts}_{int(time.time() * 1000) % 1000:03d}" + output_image = str(OUTPUT_DIR / f"{stem}.jpg") + output_json = str(OUTPUT_DIR / f"{stem}.json") + image.convert("RGB").save(output_image) + + prompt, task, output_type, target = build_prompt(ui_task, target_text, custom_prompt) + image_size = Image.open(image_path).size + stream_decoder = StreamingGeometryDecoder(image_size) + stream_drawer = GeometryImageDrawer(image_path, output_image) + stream_boxes: list[dict] = [] + stream_points: list[dict] = [] + update_queue: queue.Queue = queue.Queue() + + yield "", load_ui_image(output_image), None + + def on_token(token_id: int, step: int, piece: str, elapsed_s: float) -> None: + geometry = stream_decoder.push(token_id) + if geometry is None: + return + kind, value = geometry + if kind == "box": + assert isinstance(value, Box) + index = len(stream_boxes) + 1 + stream_drawer.add_box(value, index) + stream_boxes.append( + { + "index": index, + "token_step": step, + "elapsed_s": elapsed_s, + "box": box_to_dict(value), + } + ) + update_queue.put(("update", None, output_image, None)) + else: + assert isinstance(value, Point) + index = len(stream_points) + 1 + stream_drawer.add_point(value, index) + stream_points.append( + { + "index": index, + "token_step": step, + "elapsed_s": elapsed_s, + "point": point_to_dict(value), + } + ) + update_queue.put(("update", None, output_image, None)) + + def worker() -> None: + try: + with state.lock: + runner = state.get_runner() + token_ids, text, decoded_image_size, used_seed, timings = runner.generate( + prompt=prompt, + image_path=image_path, + max_new_tokens=int(max_new_tokens), + system_prompt=args.system_prompt, + temperature=float(temperature), + top_p=float(top_p), + repetition_penalty=float(repetition_penalty), + seed=int(seed), + on_token=on_token, + ) + + boxes, points, geometry_decode_s = time_geometry_decode(token_ids, decoded_image_size) + timings["geometry_decode_s"] = geometry_decode_s + stream_consistent = [item["box"] for item in stream_boxes] == [box_to_dict(box) for box in boxes] + stream_points_consistent = [item["point"] for item in stream_points] == [ + point_to_dict(point) for point in points + ] + + if not stream_consistent or not stream_points_consistent: + from infer_locateanything_axengine import draw_geometries + + draw_geometries(image_path, boxes, points, output_image) + + result = make_json_result( + task=task, + output_type=output_type, + target=target, + prompt=prompt, + token_ids=token_ids, + text=text, + boxes=boxes, + points=points, + stream_boxes=stream_boxes, + stream_points=stream_points, + stream_consistent=stream_consistent, + stream_points_consistent=stream_points_consistent, + timings=timings, + output_image=output_image, + ) + with open(output_json, "w", encoding="utf-8") as f: + json.dump(result, f, ensure_ascii=False, indent=2) + + ref_labels = extract_ref_labels(text) if task == "scene_text_detection" else None + result_text = format_detection_results(target, boxes, points, ref_labels) + drain_update_queue(update_queue) + update_queue.put(("final", result_text, output_image, output_json)) + except Exception as exc: + drain_update_queue(update_queue) + update_queue.put(("error", f"[ERROR] {type(exc).__name__}: {exc}", output_image, None)) + finally: + update_queue.put(("done", None, None, None)) + + thread = threading.Thread(target=worker, daemon=True) + thread.start() + + while True: + kind, log_value, image_value, file_value = update_queue.get() + if kind == "done": + break + if kind == "update": + yield gr.update(), load_ui_image(image_value), gr.update() + else: + yield log_value, load_ui_image(image_value), file_value + + return infer + + +def build_ui(state: RunnerState, args: argparse.Namespace): + infer_fn = create_infer_fn(state, args) + with gr.Blocks(title="LocateAnything") as demo: + gr.Markdown("# LocateAnything Gradio UI") + gr.Markdown( + "上传图片，选择任务类型，输入类别、category set 或 phrase 后开始推理。OCR 任务可留空；box/point 会在生成过程中逐个写入输出图。" + ) + with gr.Row(): + with gr.Column(scale=1): + image = gr.Image(label="图片上传", type="pil") + task = gr.Dropdown( + label="任务类别", + choices=list(UI_TASKS), + value="Object Detection", + ) + target = gr.Textbox(label="类别 / category set / phrase（OCR 可留空）", value=DEFAULT_TARGET) + custom_prompt = gr.Textbox(label="自定义 prompt（可选，填写后覆盖任务模板）", value="") + with gr.Accordion("生成参数", open=False): + max_new_tokens = gr.Slider(32, 512, value=512, step=1, label="max_new_tokens") + temperature = gr.Slider(0.0, 2.0, value=0.7, step=0.05, label="temperature") + top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.01, label="top_p") + repetition_penalty = gr.Slider(1.0, 2.0, value=1.1, step=0.05, label="repetition_penalty") + seed = gr.Number(value=42, precision=0, label="seed") + run_btn = gr.Button("开始推理", variant="primary") + with gr.Column(scale=1): + output_image = gr.Image(label="原图结果流式可视化", type="pil") + log = gr.Textbox(label="最终检测结果 [ref, box]", lines=24) + json_file = gr.File(label="标准 JSON 结果下载") + + run_btn.click( + infer_fn, + inputs=[image, task, target, custom_prompt, max_new_tokens, temperature, top_p, repetition_penalty, seed], + outputs=[log, output_image, json_file], + ) + return demo + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="LocateAnything Gradio UI") + parser.add_argument("--tokenizer", default=DEFAULT_TOKENIZER) + parser.add_argument("--llm-dir", default=DEFAULT_LLM_DIR) + parser.add_argument("--image-encoder", default=DEFAULT_IMAGE_ENCODER) + parser.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT) + parser.add_argument("--host", default="0.0.0.0") + parser.add_argument("--port", type=int, default=7860) + parser.add_argument("--port-retries", type=int, default=20, help="Try following ports when --port is occupied.") + parser.add_argument("--share", action="store_true") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + ensure_gradio() + UI_DIR.mkdir(parents=True, exist_ok=True) + state = RunnerState(args.tokenizer, args.llm_dir, args.image_encoder) + demo = build_ui(state, args) + selected_port = find_available_port(args.port, args.port_retries) + if selected_port != args.port: + print(f"[Gradio] port {args.port} is occupied, fallback to {selected_port}", flush=True) + + print("[Gradio] starting LocateAnything UI", flush=True) + print(f"[Gradio] local: http://127.0.0.1:{selected_port}", flush=True) + for ip in get_lan_ips(): + print(f"[Gradio] LAN: http://{ip}:{selected_port}", flush=True) + print("[Gradio] Use another computer in the same LAN to open the LAN URL.", flush=True) + + try: + demo.queue(default_concurrency_limit=1).launch( + server_name=args.host, + server_port=selected_port, + share=args.share, + ) + finally: + state.close() + + +if __name__ == "__main__": + main() diff --git a/image_encoder_mlp.axmodel b/image_encoder_mlp.axmodel new file mode 100644 index 0000000000000000000000000000000000000000..cbfeeb82600a336a4617e955698b9f6411c760a1 --- /dev/null +++ b/image_encoder_mlp.axmodel @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6f72727633e605bf7eda5eb0fd478a425297e065b161b9029f027c89cb545f93 +size 592696349 diff --git a/infer_locateanything_axengine.py b/infer_locateanything_axengine.py new file mode 100644 index 0000000000000000000000000000000000000000..5d91c2b6974734c7e66b27680414ad7b50f640c9 --- /dev/null +++ b/infer_locateanything_axengine.py @@ -0,0 +1,997 @@ +#!/usr/bin/env python3 +"""Pure Python LocateAnything inference with axengine. + +This runner does not use axllm serve or OpenAI-compatible HTTP APIs. It runs: + 1. image_encoder_mlp.axmodel + 2. llm-4bit-650/qwen2_p128_l*_together.axmodel + 3. llm-4bit-650/qwen2_post.axmodel + +The implementation mirrors infer_axmodel.py's prefill/decode/post flow and +injects image encoder embeddings at LocateAnything's positions. +""" + +from __future__ import annotations + +import argparse +import atexit +import json +import os +import time +from dataclasses import dataclass +from typing import Sequence + +import numpy as np +from axengine import InferenceSession +from ml_dtypes import bfloat16 +from PIL import Image, ImageDraw, ImageFont +from tokenizers import Tokenizer + + +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +DEFAULT_TOKENIZER = os.path.join(SCRIPT_DIR, "qwen2.5_tokenizer", "tokenizer.json") +DEFAULT_LLM_DIR = os.path.join(SCRIPT_DIR) +DEFAULT_IMAGE_ENCODER = os.path.join(SCRIPT_DIR, "image_encoder_mlp.axmodel") +DEFAULT_IMAGE = os.path.join(SCRIPT_DIR, "test_data", "person.jpg") +DEFAULT_SYSTEM_PROMPT = "You are a helpful assistant." +DEFAULT_TASK = "object_detection" +DEFAULT_TARGET = "person" + +PATCH_SIZE = 14 +VISION_SIZE = 560 +VISION_TOKENS = 400 +HIDDEN_SIZE = 2048 +VOCAB_SIZE = 152681 +KV_CACHE_LEN = 1024 + +TOK_BOX_START = 151668 +TOK_BOX_END = 151669 +TOK_COORD_START = 151677 +TOK_COORD_END = 152677 +TOK_IMG_CONTEXT = 151665 +TOK_EOS = 151645 + + +@dataclass +class Box: + x1: float + y1: float + x2: float + y2: float + + +@dataclass +class Point: + x: float + y: float + + +@dataclass +class StreamBox: + index: int + token_step: int + elapsed_s: float + box: Box + + +@dataclass +class StreamPoint: + index: int + token_step: int + elapsed_s: float + point: Point + + +@dataclass +class LayerFiles: + layer_paths: list[str] + post_path: str + embed_path: str + + +@dataclass(frozen=True) +class PromptSpec: + output_type: str + target_kind: str | None + template: str + + +PROMPT_SPECS = { + "object_detection": PromptSpec( + "box", + "categories", + "Locate all the instances that matches the following description:{target}", + ), + "phrase_grounding_single": PromptSpec( + "box", + "phrase", + "Locate a single instance that matches the following description: {target}.", + ), + "phrase_grounding_multi": PromptSpec( + "box", + "phrase", + "Locate all the instances that match the following description: {target}.", + ), + "text_grounding": PromptSpec("box", "phrase", "Please locate the text referred as {target}."), + "scene_text_detection": PromptSpec("box", None, "Detect all the text in box format."), + "document_layout": PromptSpec( + "box", + "categories", + "Detect all the objects in the image that belong to the category set: {target}.", + ), + "gui_grounding_box": PromptSpec( + "box", + "phrase", + "Locate the region that matches the following description: {target}.", + ), + "gui_grounding_point": PromptSpec("point", "phrase", "Point to: {target}."), + "pointing": PromptSpec("point", "phrase", "Point to: {target}."), +} + + +def release_ax_inference_session(session) -> None: + inner = getattr(session, "_sess", None) + unload = getattr(inner, "_unload", None) + if not callable(unload): + return + try: + unload() + except Exception as exc: + print(f"[WARN] Failed to unload axengine session cleanly: {exc}") + finally: + try: + inner._unload = lambda: None + except Exception: + pass + + +def dtype_from_axengine(dtype) -> np.dtype: + name = str(dtype).lower() + if "bfloat16" in name or "bf16" in name: + return bfloat16 + if "float32" in name or "fp32" in name: + return np.float32 + if "float16" in name or "fp16" in name: + return np.float16 + if "uint32" in name or "u32" in name: + return np.uint32 + raise ValueError(f"Unsupported axengine dtype: {dtype}") + + +def ensure_finite(name: str, arr: np.ndarray) -> None: + arr32 = np.asarray(arr, dtype=np.float32) + if not np.isfinite(arr32).all(): + raise RuntimeError(f"{name} contains NaN/Inf") + + +def resolve_model_files(llm_dir: str) -> LayerFiles: + cfg_path = os.path.join(llm_dir, "config.json") + with open(cfg_path, "r", encoding="utf-8") as f: + cfg = json.load(f) + + template = cfg.get("template_filename_axmodel", "qwen2_p128_l%d_together.axmodel") + num_layers = int(cfg.get("axmodel_num", 36)) + layer_paths = [os.path.join(llm_dir, template % i) for i in range(num_layers)] + missing = [p for p in layer_paths if not os.path.exists(p)] + if missing: + raise FileNotFoundError(f"Missing layer axmodel: {missing[0]}") + + post_path = os.path.join(llm_dir, cfg.get("filename_post_axmodel", "qwen2_post.axmodel")) + embed_path = os.path.join(llm_dir, cfg.get("filename_tokens_embed", "model.embed_tokens.weight.bfloat16.bin")) + if not os.path.exists(post_path): + raise FileNotFoundError(post_path) + if not os.path.exists(embed_path): + raise FileNotFoundError(embed_path) + return LayerFiles(layer_paths=layer_paths, post_path=post_path, embed_path=embed_path) + + +def image_to_patches(image_path: str) -> tuple[np.ndarray, tuple[int, int]]: + image = Image.open(image_path).convert("RGB") + original_size = image.size + if image.size != (VISION_SIZE, VISION_SIZE): + image = image.resize((VISION_SIZE, VISION_SIZE), Image.Resampling.BICUBIC) + + # Match C++ LocateAnythingImageProcessor: uint8 RGB after Pillow-like + # resize, then normalize each pixel as v / 127.5 - 1.0. + arr = np.asarray(image, dtype=np.uint8).astype(np.float32) + arr = arr / 127.5 - 1.0 + arr = arr.transpose(2, 0, 1) + + c, h, w = arr.shape + h_grid = h // PATCH_SIZE + w_grid = w // PATCH_SIZE + patches = arr.reshape(c, h_grid, PATCH_SIZE, w_grid, PATCH_SIZE) + patches = patches.transpose(1, 3, 0, 2, 4) + patches = np.ascontiguousarray(patches.reshape(-1, c, PATCH_SIZE, PATCH_SIZE), dtype=np.float32) + return patches, original_size + + +def build_locateanything_prompt( + system_prompt: str, + prompt: str, + image_tokens: int = VISION_TOKENS, +) -> str: + media = "" + "" * image_tokens + "" + return ( + f"<|im_start|>system\n{system_prompt}\n<|im_end|>\n" + "<|im_start|>user\n" + + media + + prompt + + "<|im_end|>\n<|im_start|>assistant\n" + ) + + +def make_feed(shapes: dict[str, tuple[int, ...]], values: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + return {name: value for name, value in values.items() if name in shapes} + + +def build_task_prompt(args: argparse.Namespace) -> tuple[str, str, str, str]: + if args.prompt: + return args.prompt, "custom", "unknown", "custom" + + spec = PROMPT_SPECS[args.task] + if spec.target_kind == "categories": + target = args.categories or args.target + elif spec.target_kind == "phrase": + target = args.phrase or args.target + else: + target = "" + + if spec.target_kind is not None and not target: + raise ValueError(f"--task {args.task} requires --target, --categories, or --phrase") + + prompt = spec.template.format(target=target) + return prompt, args.task, spec.output_type, target + + +class LocateAnythingAxEngineRunner: + def __init__(self, tokenizer_path: str, llm_dir: str, image_encoder_path: str): + self.tokenizer = Tokenizer.from_file(tokenizer_path) + self.files = resolve_model_files(llm_dir) + self.embed_matrix = np.memmap(self.files.embed_path, mode="r", dtype=np.uint16).view(bfloat16).reshape( + VOCAB_SIZE, HIDDEN_SIZE + ) + + self.image_encoder = InferenceSession(image_encoder_path) + self.layer_sessions = [InferenceSession(path) for path in self.files.layer_paths] + self.post_session = InferenceSession(self.files.post_path) + self._closed = False + atexit.register(self.close) + + self.layer_decode_input_shapes: list[dict[str, tuple[int, ...]]] = [] + self.layer_decode_input_dtypes: list[dict[str, np.dtype]] = [] + self.layer_decode_output_names: list[list[str]] = [] + self.layer_prefill_input_shapes: list[list[dict[str, tuple[int, ...]]]] = [] + self.layer_prefill_input_dtypes: list[list[dict[str, np.dtype]]] = [] + self.layer_prefill_output_names: list[list[list[str]]] = [] + + for session in self.layer_sessions: + decode_inputs = session.get_inputs(shape_group=0) + self.layer_decode_input_shapes.append({x.name: tuple(x.shape) for x in decode_inputs}) + self.layer_decode_input_dtypes.append({x.name: dtype_from_axengine(x.dtype) for x in decode_inputs}) + self.layer_decode_output_names.append([x.name for x in session.get_outputs(shape_group=0)]) + + group_shapes = [] + group_dtypes = [] + group_outputs = [] + for shape_group in range(1, 64): + try: + inputs = session.get_inputs(shape_group=shape_group) + outputs = session.get_outputs(shape_group=shape_group) + except Exception: + break + group_shapes.append({x.name: tuple(x.shape) for x in inputs}) + group_dtypes.append({x.name: dtype_from_axengine(x.dtype) for x in inputs}) + group_outputs.append([x.name for x in outputs]) + self.layer_prefill_input_shapes.append(group_shapes) + self.layer_prefill_input_dtypes.append(group_dtypes) + self.layer_prefill_output_names.append(group_outputs) + + self.prefill_len = int(self.layer_prefill_input_shapes[0][0]["input"][1]) + self.hidden_dtype = self.layer_decode_input_dtypes[0]["input"] + + def close(self) -> None: + if self._closed: + return + release_ax_inference_session(self.image_encoder) + for session in self.layer_sessions: + release_ax_inference_session(session) + release_ax_inference_session(self.post_session) + self._closed = True + + def encode_text(self, text: str) -> list[int]: + return self.tokenizer.encode(text, add_special_tokens=False).ids + + def decode_text(self, token_ids: Sequence[int]) -> str: + return self.tokenizer.decode(list(token_ids), skip_special_tokens=False) + + def embed_token(self, token_id: int) -> np.ndarray: + return np.asarray(self.embed_matrix[int(token_id)], dtype=self.hidden_dtype).reshape(1, 1, HIDDEN_SIZE) + + def alloc_layer_states(self) -> tuple[list[np.ndarray], list[np.ndarray]]: + k_states = [] + v_states = [] + for shapes, dtypes in zip(self.layer_decode_input_shapes, self.layer_decode_input_dtypes): + k_states.append(np.zeros(shapes["K_cache"], dtype=dtypes["K_cache"])) + v_states.append(np.zeros(shapes["V_cache"], dtype=dtypes["V_cache"])) + return k_states, v_states + + def encode_image(self, image_path: str) -> tuple[np.ndarray, tuple[int, int], dict[str, float]]: + t0 = time.perf_counter() + patches, original_size = image_to_patches(image_path) + preprocess_s = time.perf_counter() - t0 + + t1 = time.perf_counter() + input_name = self.image_encoder.get_inputs()[0].name + output = self.image_encoder.run(None, {input_name: patches})[0] + image_encoder_s = time.perf_counter() - t1 + + output = np.asarray(output, dtype=np.float32) + if output.shape != (VISION_TOKENS, HIDDEN_SIZE): + raise RuntimeError(f"Unexpected image encoder output shape: {output.shape}") + ensure_finite("image encoder output", output) + timings = { + "image_preprocess_s": preprocess_s, + "image_encoder_s": image_encoder_s, + "image_total_s": preprocess_s + image_encoder_s, + } + return output.astype(self.hidden_dtype), original_size, timings + + def build_prompt_embeddings( + self, + prompt: str, + image_embeds: np.ndarray, + system_prompt: str, + ) -> tuple[list[int], np.ndarray]: + text = build_locateanything_prompt( + system_prompt, + prompt, + image_tokens=image_embeds.shape[0], + ) + token_ids = self.encode_text(text) + image_positions = [i for i, token_id in enumerate(token_ids) if token_id == TOK_IMG_CONTEXT] + if len(image_positions) != image_embeds.shape[0]: + raise RuntimeError( + f"image placeholder count mismatch: prompt={len(image_positions)} image_embeds={image_embeds.shape[0]}" + ) + + embeds = np.asarray(self.embed_matrix[np.asarray(token_ids, dtype=np.int64)], dtype=self.hidden_dtype) + embeds[image_positions, :] = image_embeds + return token_ids, embeds + + @staticmethod + def prefill_history_capacity(shapes: dict[str, tuple[int, ...]]) -> int: + input_len = int(shapes.get("input", (1, 128))[1]) + mask_shape = shapes.get("mask") + if mask_shape is not None and len(mask_shape) == 3: + return max(0, int(mask_shape[-1]) - input_len) + k_shape = shapes.get("K_cache") + if k_shape is not None and len(k_shape) >= 2: + return int(k_shape[1]) + return 0 + + def select_prefill_shape_group(self, layer_idx: int, history_len: int) -> int: + groups = self.layer_prefill_input_shapes[layer_idx] + candidates = [] + for offset, shapes in enumerate(groups): + cap = self.prefill_history_capacity(shapes) + if cap >= history_len: + candidates.append((cap, offset + 1)) + if candidates: + return min(candidates)[1] + return max(range(1, len(groups) + 1), key=lambda gid: self.prefill_history_capacity(groups[gid - 1])) + + def run_prefill(self, prompt_embeds: np.ndarray): + k_states, v_states = self.alloc_layer_states() + last_hidden = None + total_len = int(prompt_embeds.shape[0]) + + for start in range(0, total_len, self.prefill_len): + chunk = prompt_embeds[start : start + self.prefill_len] + chunk_len = int(chunk.shape[0]) + data = np.zeros((1, self.prefill_len, HIDDEN_SIZE), dtype=self.hidden_dtype) + data[0, :chunk_len, :] = chunk + + for layer_idx, session in enumerate(self.layer_sessions): + shape_group = self.select_prefill_shape_group(layer_idx, start) + layer_shapes = self.layer_prefill_input_shapes[layer_idx][shape_group - 1] + layer_dtypes = self.layer_prefill_input_dtypes[layer_idx][shape_group - 1] + history_cap = self.prefill_history_capacity(layer_shapes) + history_len = min(start, history_cap) + + indices = np.zeros(layer_shapes["indices"], dtype=layer_dtypes["indices"]) + indices.reshape(-1)[:chunk_len] = np.arange(start, start + chunk_len, dtype=np.uint32) + + mask = np.full(layer_shapes["mask"], -65536.0, dtype=np.float32) + for q in range(chunk_len): + mask[:, q, : history_len + q + 1] = 0.0 + mask = mask.astype(layer_dtypes["mask"]) + + k_feed = np.zeros(layer_shapes["K_cache"], dtype=layer_dtypes["K_cache"]) + v_feed = np.zeros(layer_shapes["V_cache"], dtype=layer_dtypes["V_cache"]) + if history_len > 0: + k_feed[:, :history_len, :] = k_states[layer_idx][:, :history_len, :] + v_feed[:, :history_len, :] = v_states[layer_idx][:, :history_len, :] + + outputs = session.run( + None, + make_feed( + layer_shapes, + { + "K_cache": k_feed, + "V_cache": v_feed, + "indices": indices, + "input": data.astype(layer_dtypes["input"], copy=False), + "mask": mask, + }, + ), + shape_group=shape_group, + ) + output_map = dict(zip(self.layer_prefill_output_names[layer_idx][shape_group - 1], outputs)) + k_out = output_map.get("K_cache_out") + if k_out is not None: + k_states[layer_idx][:, start : start + chunk_len, :] = k_out[:, :chunk_len, :] + v_out = output_map.get("V_cache_out") + if v_out is not None: + v_states[layer_idx][:, start : start + chunk_len, :] = v_out[:, :chunk_len, :] + data = output_map["output"] + ensure_finite(f"prefill layer {layer_idx}", data) + + last_hidden = data[:, chunk_len - 1 : chunk_len, :] + return k_states, v_states, last_hidden + + def run_decode_step( + self, + hidden: np.ndarray, + position: int, + k_states: list[np.ndarray], + v_states: list[np.ndarray], + ) -> np.ndarray: + data = hidden + for layer_idx, session in enumerate(self.layer_sessions): + layer_shapes = self.layer_decode_input_shapes[layer_idx] + layer_dtypes = self.layer_decode_input_dtypes[layer_idx] + + indices = np.zeros(layer_shapes["indices"], dtype=layer_dtypes["indices"]) + indices.reshape(-1)[0] = position + + mask = np.full(layer_shapes["mask"], -65536.0, dtype=np.float32) + valid_past = min(position, layer_shapes["mask"][-1] - 1) + if valid_past > 0: + mask[:, :, :valid_past] = 0.0 + mask[:, :, -1:] = 0.0 + mask = mask.astype(layer_dtypes["mask"]) + + outputs = session.run( + None, + make_feed( + layer_shapes, + { + "K_cache": k_states[layer_idx], + "V_cache": v_states[layer_idx], + "indices": indices, + "input": data.astype(layer_dtypes["input"], copy=False), + "mask": mask, + }, + ), + shape_group=0, + ) + output_map = dict(zip(self.layer_decode_output_names[layer_idx], outputs)) + k_out = output_map.get("K_cache_out") + if k_out is not None: + pos_end = position + k_out.shape[1] + k_states[layer_idx][:, position:pos_end, :] = k_out + v_out = output_map.get("V_cache_out") + if v_out is not None: + pos_end = position + v_out.shape[1] + v_states[layer_idx][:, position:pos_end, :] = v_out + data = output_map["output"] + ensure_finite(f"decode layer {layer_idx}", data) + return data + + def run_post(self, hidden: np.ndarray) -> np.ndarray: + logits = self.post_session.run(None, {"input": hidden.astype(self.hidden_dtype, copy=False)})[0] + ensure_finite("post logits", logits) + return np.asarray(logits, dtype=np.float32).reshape(-1) + + @staticmethod + def sample_next_token( + logits: np.ndarray, + rng: np.random.Generator | None, + temperature: float, + top_p: float, + repetition_penalty: float, + generated_history: Sequence[int], + ) -> int: + scores = logits.astype(np.float64) + if repetition_penalty != 1.0 and generated_history: + for token_id in set(int(x) for x in generated_history if 0 <= int(x) < scores.shape[0]): + if scores[token_id] > 0: + scores[token_id] /= repetition_penalty + else: + scores[token_id] *= repetition_penalty + + if temperature > 0.0: + scores = scores / float(temperature) + if top_p is not None and top_p < 1.0: + order = np.argsort(scores)[::-1] + sorted_scores = scores[order] + stable_scores = sorted_scores - np.max(sorted_scores) + sorted_probs = np.exp(stable_scores) + sorted_probs /= np.sum(sorted_probs) + cumulative = np.cumsum(sorted_probs) + remove = cumulative > top_p + if remove.shape[0] > 1: + remove[1:] = remove[:-1] + remove[0] = False + scores[order[remove]] = -np.inf + + if temperature <= 0.0: + return int(np.argmax(scores)) + + scores -= np.max(scores) + probs = np.exp(scores) + probs_sum = float(np.sum(probs)) + if not np.isfinite(probs_sum) or probs_sum <= 0.0: + return int(np.argmax(logits)) + probs /= probs_sum + + assert rng is not None + return int(rng.choice(np.arange(probs.shape[0]), p=probs)) + + def generate_from_state( + self, + prompt_len: int, + k_states: list[np.ndarray], + v_states: list[np.ndarray], + last_hidden: np.ndarray, + max_new_tokens: int, + temperature: float, + top_p: float, + repetition_penalty: float, + seed: int, + prompt_token_ids: Sequence[int], + on_token=None, + ) -> list[int]: + rng = None if temperature <= 0.0 else np.random.default_rng(seed) + generated: list[int] = [] + generated_history = list(prompt_token_ids) + decode_start = time.perf_counter() + for step in range(max_new_tokens): + next_token = self.sample_next_token( + self.run_post(last_hidden), + rng, + temperature, + top_p, + repetition_penalty, + generated_history, + ) + generated.append(next_token) + generated_history.append(next_token) + if on_token is not None: + on_token( + next_token, + step, + self.decode_text([next_token]), + time.perf_counter() - decode_start, + ) + if next_token == TOK_EOS: + break + last_hidden = self.run_decode_step( + self.embed_token(next_token), + prompt_len + step, + k_states, + v_states, + ) + return generated + + def generate( + self, + prompt: str, + image_path: str, + max_new_tokens: int, + system_prompt: str, + temperature: float, + top_p: float, + repetition_penalty: float, + seed: int, + on_token=None, + ) -> tuple[list[int], str, tuple[int, int], int, dict[str, float]]: + image_embeds, image_size, timings = self.encode_image(image_path) + print( + f"[Image] encoded {image_path}, embeds={image_embeds.shape}, " + f"preprocess={timings['image_preprocess_s']:.3f}s " + f"encoder={timings['image_encoder_s']:.3f}s " + f"total={timings['image_total_s']:.3f}s" + ) + + t_prompt = time.perf_counter() + token_ids, prompt_embeds = self.build_prompt_embeddings(prompt, image_embeds, system_prompt) + timings["prompt_build_s"] = time.perf_counter() - t_prompt + print(f"[Prompt] tokens={len(token_ids)}, image_tokens={token_ids.count(TOK_IMG_CONTEXT)} style=original") + if len(token_ids) + max_new_tokens >= KV_CACHE_LEN: + raise RuntimeError(f"prompt + max_new_tokens exceeds KV cache: {len(token_ids)} + {max_new_tokens} >= {KV_CACHE_LEN}") + + t1 = time.perf_counter() + k_states, v_states, last_hidden = self.run_prefill(prompt_embeds) + timings["llm_prefill_s"] = time.perf_counter() - t1 + print(f"[Prefill] cost={timings['llm_prefill_s']:.3f}s") + + t2 = time.perf_counter() + print( + f"[LLM Decode] seed={seed} temperature={temperature} " + f"top_p={top_p} repetition_penalty={repetition_penalty}" + ) + generated = self.generate_from_state( + len(token_ids), + k_states, + v_states, + last_hidden, + max_new_tokens, + temperature, + top_p, + repetition_penalty, + seed, + token_ids, + on_token=on_token, + ) + timings["llm_generate_s"] = time.perf_counter() - t2 + timings["llm_total_s"] = timings["llm_prefill_s"] + timings["llm_generate_s"] + print(f"[LLM Decode] generated={len(generated)} cost={timings['llm_generate_s']:.3f}s") + print(f"[Timing] llm_total={timings['llm_total_s']:.3f}s") + return generated, self.decode_text(generated), image_size, seed, timings + + +def token_to_coord(token_id: int) -> int | None: + if TOK_COORD_START <= token_id <= TOK_COORD_END: + return token_id - TOK_COORD_START + return None + + +def coord_to_pixel(coord: int, dim: int) -> float: + return coord / 1000 * dim + + +def make_box(coords: Sequence[int], image_size: tuple[int, int]) -> Box: + iw, ih = image_size + x1, y1, x2, y2 = [int(v) for v in coords] + return Box( + coord_to_pixel(x1, iw), + coord_to_pixel(y1, ih), + coord_to_pixel(x2, iw), + coord_to_pixel(y2, ih), + ) + + +def make_point(coords: Sequence[int], image_size: tuple[int, int]) -> Point: + iw, ih = image_size + x, y = [int(v) for v in coords] + return Point(coord_to_pixel(x, iw), coord_to_pixel(y, ih)) + + +def decode_boxes_strict(token_ids: Sequence[int], image_size: tuple[int, int]) -> list[Box]: + """Decode exact token sequences.""" + boxes: list[Box] = [] + i = 0 + while i <= len(token_ids) - 6: + if token_ids[i] != TOK_BOX_START: + i += 1 + continue + coords = [token_to_coord(token_ids[i + j]) for j in range(1, 5)] + if None not in coords and token_ids[i + 5] == TOK_BOX_END: + boxes.append(make_box([int(v) for v in coords], image_size)) + i += 6 + else: + i += 1 + return boxes + + +def decode_points_strict(token_ids: Sequence[int], image_size: tuple[int, int]) -> list[Point]: + """Decode exact point token sequences.""" + points: list[Point] = [] + i = 0 + while i <= len(token_ids) - 4: + if token_ids[i] != TOK_BOX_START: + i += 1 + continue + coords = [token_to_coord(token_ids[i + 1]), token_to_coord(token_ids[i + 2])] + if None not in coords and token_ids[i + 3] == TOK_BOX_END: + points.append(make_point([int(v) for v in coords], image_size)) + i += 4 + else: + i += 1 + return points + + +class StreamingGeometryDecoder: + """Incrementally decode strict box and point token sequences.""" + + def __init__(self, image_size: tuple[int, int]): + self.image_size = image_size + self._state = "search" + self._coords: list[int] = [] + + def push(self, token_id: int) -> tuple[str, Box | Point] | None: + if self._state == "search": + if token_id == TOK_BOX_START: + self._coords = [] + self._state = "coords" + return None + + if self._state == "coords": + coord = token_to_coord(token_id) + if coord is not None: + self._coords.append(coord) + if len(self._coords) == 4: + self._state = "end" + return None + if token_id == TOK_BOX_END and len(self._coords) == 2: + point = make_point(self._coords, self.image_size) + self._state = "search" + self._coords = [] + return "point", point + self._reset(token_id) + return None + + if self._state == "end": + if token_id == TOK_BOX_END: + box = make_box(self._coords, self.image_size) + self._state = "search" + self._coords = [] + return "box", box + self._reset(token_id) + return None + + def _reset(self, token_id: int) -> None: + if token_id == TOK_BOX_START: + self._coords = [] + self._state = "coords" + else: + self._coords = [] + self._state = "search" + + +def time_geometry_decode( + token_ids: Sequence[int], + image_size: tuple[int, int], +) -> tuple[list[Box], list[Point], float]: + t0 = time.perf_counter() + boxes = decode_boxes_strict(token_ids, image_size) + points = decode_points_strict(token_ids, image_size) + return boxes, points, time.perf_counter() - t0 + + +BOX_COLORS = ["#00FF00", "#FF0000", "#0000FF", "#FFFF00", "#FF00FF", "#00FFFF", "#FFA500"] + + +def load_box_font(): + try: + return ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 16) + except OSError: + return ImageFont.load_default() + + +def draw_one_box(draw: ImageDraw.ImageDraw, font, box: Box, idx: int) -> None: + color = BOX_COLORS[(idx - 1) % len(BOX_COLORS)] + rect = [round(box.x1), round(box.y1), round(box.x2), round(box.y2)] + draw.rectangle(rect, outline=color, width=3) + label = str(idx) + label_pos = (rect[0], max(0, rect[1] - 18)) + tb = draw.textbbox(label_pos, label, font=font) + draw.rectangle(tb, fill=color) + draw.text(label_pos, label, fill="black", font=font) + + +def draw_one_point(draw: ImageDraw.ImageDraw, font, point: Point, idx: int) -> None: + color = BOX_COLORS[(idx - 1) % len(BOX_COLORS)] + x = round(point.x) + y = round(point.y) + radius = 8 + draw.ellipse([x - radius, y - radius, x + radius, y + radius], outline=color, fill=color, width=3) + label = str(idx) + label_pos = (x + radius + 2, max(0, y - radius)) + tb = draw.textbbox(label_pos, label, font=font) + draw.rectangle(tb, fill=color) + draw.text(label_pos, label, fill="black", font=font) + + +def draw_geometries(image_path: str, boxes: list[Box], points: list[Point], output: str) -> None: + image = Image.open(image_path).convert("RGB") + draw = ImageDraw.Draw(image) + font = load_box_font() + for idx, box in enumerate(boxes, start=1): + draw_one_box(draw, font, box, idx) + for idx, point in enumerate(points, start=1): + draw_one_point(draw, font, point, idx) + save_image_atomic(image, output) + + +def save_image_atomic(image: Image.Image, output: str) -> None: + output_dir = os.path.dirname(os.path.abspath(output)) + os.makedirs(output_dir, exist_ok=True) + root, ext = os.path.splitext(output) + tmp_output = f"{root}.tmp.{os.getpid()}.{time.time_ns()}{ext or '.png'}" + try: + image.save(tmp_output) + os.replace(tmp_output, output) + finally: + if os.path.exists(tmp_output): + os.unlink(tmp_output) + + +class GeometryImageDrawer: + def __init__(self, image_path: str, output: str): + self.output = output + output_dir = os.path.dirname(os.path.abspath(output)) + os.makedirs(output_dir, exist_ok=True) + self.image = Image.open(image_path).convert("RGB") + self.draw = ImageDraw.Draw(self.image) + self.font = load_box_font() + + def add_box(self, box: Box, idx: int) -> None: + draw_one_box(self.draw, self.font, box, idx) + save_image_atomic(self.image, self.output) + + def add_point(self, point: Point, idx: int) -> None: + draw_one_point(self.draw, self.font, point, idx) + save_image_atomic(self.image, self.output) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="LocateAnything pure axengine Python inference") + parser.add_argument("--tokenizer", default=DEFAULT_TOKENIZER) + parser.add_argument("--llm-dir", default=DEFAULT_LLM_DIR) + parser.add_argument("--image-encoder", default=DEFAULT_IMAGE_ENCODER) + parser.add_argument("--image", default=DEFAULT_IMAGE) + parser.add_argument("--prompt", default=None, help="Custom prompt. Overrides --task prompt templates.") + parser.add_argument("--task", choices=sorted(PROMPT_SPECS), default=DEFAULT_TASK) + parser.add_argument("--target", default=DEFAULT_TARGET, help="Fallback target text for prompt templates.") + parser.add_argument("--categories", default=None, help="Comma-separated category names for category tasks.") + parser.add_argument("--phrase", default=None, help="Free-form phrase for grounding/pointing tasks.") + parser.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT) + parser.add_argument("--output", default="output_locateanything_axengine.jpg") + parser.add_argument("--max-new-tokens", type=int, default=512) + parser.add_argument("--temperature", type=float, default=0.7) + parser.add_argument("--top-p", type=float, default=0.9) + parser.add_argument("--repetition-penalty", type=float, default=1.1) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--save-response", default=None) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + prompt, prompt_task, prompt_output_type, prompt_target = build_task_prompt(args) + stream_decoder: StreamingGeometryDecoder | None = None + stream_drawer: GeometryImageDrawer | None = None + stream_boxes: list[StreamBox] = [] + stream_points: list[StreamPoint] = [] + + def handle_stream_token(token_id: int, step: int, piece: str, elapsed_s: float) -> None: + assert stream_decoder is not None + assert stream_drawer is not None + print(f"[Stream Token] step={step:03d} token={token_id} text={piece!r}") + geometry = stream_decoder.push(token_id) + if geometry is None: + return + kind, value = geometry + if kind == "point": + assert isinstance(value, Point) + stream_point = StreamPoint( + index=len(stream_points) + 1, + token_step=step, + elapsed_s=elapsed_s, + point=value, + ) + stream_points.append(stream_point) + stream_drawer.add_point(value, stream_point.index) + print( + f"[Stream Point] #{stream_point.index} step={step} " + f"elapsed={elapsed_s:.3f}s saved={args.output}" + ) + return + + assert isinstance(value, Box) + stream_box = StreamBox( + index=len(stream_boxes) + 1, + token_step=step, + elapsed_s=elapsed_s, + box=value, + ) + stream_boxes.append(stream_box) + stream_drawer.add_box(value, stream_box.index) + print( + f"[Stream Box] #{stream_box.index} step={step} " + f"elapsed={elapsed_s:.3f}s saved={args.output}" + ) + + print( + f"[Task Prompt] task={prompt_task} output={prompt_output_type} " + f"target={prompt_target!r} prompt={prompt!r}" + ) + runner = LocateAnythingAxEngineRunner(args.tokenizer, args.llm_dir, args.image_encoder) + try: + image_size_hint = Image.open(args.image).size + stream_decoder = StreamingGeometryDecoder(image_size_hint) + stream_drawer = GeometryImageDrawer(args.image, args.output) + token_ids, text, image_size, used_seed, timings = runner.generate( + prompt=prompt, + image_path=args.image, + max_new_tokens=args.max_new_tokens, + system_prompt=args.system_prompt, + temperature=args.temperature, + top_p=args.top_p, + repetition_penalty=args.repetition_penalty, + seed=args.seed, + on_token=handle_stream_token, + ) + finally: + runner.close() + + boxes, points, decode_s = time_geometry_decode(token_ids, image_size) + timings["geometry_decode_s"] = decode_s + stream_consistent = [x.box for x in stream_boxes] == boxes + stream_points_consistent = [x.point for x in stream_points] == points + if not stream_consistent or not stream_points_consistent: + print("[WARN] Stream geometries differ from final strict decode; rewriting final image.") + draw_geometries(args.image, boxes, points, args.output) + + print("\n[LLM output]") + print(text) + print(f"\n[Seed] {used_seed}") + print(f"\n[Geometries] boxes={len(boxes)} points={len(points)} decode={decode_s * 1000:.3f}ms") + print( + f"[Stream] boxes={len(stream_boxes)} consistent={stream_consistent} " + f"points={len(stream_points)} points_consistent={stream_points_consistent}" + ) + for i, box in enumerate(boxes, start=1): + print(f" [{i}] ({box.x1:.2f},{box.y1:.2f}) -> ({box.x2:.2f},{box.y2:.2f})") + for i, point in enumerate(points, start=1): + print(f" [P{i}] ({point.x:.2f},{point.y:.2f})") + + if args.save_response: + with open(args.save_response, "w", encoding="utf-8") as f: + json.dump( + { + "task": prompt_task, + "task_output_type": prompt_output_type, + "target": prompt_target, + "prompt": prompt, + "seed": used_seed, + "temperature": args.temperature, + "top_p": args.top_p, + "repetition_penalty": args.repetition_penalty, + "token_ids": token_ids, + "text": text, + "boxes": [box.__dict__ for box in boxes], + "points": [point.__dict__ for point in points], + "stream_boxes": [ + { + "index": x.index, + "token_step": x.token_step, + "elapsed_s": x.elapsed_s, + "box": x.box.__dict__, + } + for x in stream_boxes + ], + "stream_points": [ + { + "index": x.index, + "token_step": x.token_step, + "elapsed_s": x.elapsed_s, + "point": x.point.__dict__, + } + for x in stream_points + ], + "stream_consistent": stream_consistent, + "stream_points_consistent": stream_points_consistent, + "timings": timings, + }, + f, + ensure_ascii=False, + ) + print("[Save] response:", args.save_response) + print("[Draw] saved:", args.output) + + +if __name__ == "__main__": + main() diff --git a/model.embed_tokens.weight.bfloat16.bin b/model.embed_tokens.weight.bfloat16.bin new file mode 100644 index 0000000000000000000000000000000000000000..099cea0f449327062e3737704333735fe9d1fe10 --- /dev/null +++ b/model.embed_tokens.weight.bfloat16.bin @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7972da8526034a169b9a0f7e5478d3fe2a5ce61e9dd2a74b3a5b9ad1e4759db9 +size 625381376 diff --git a/post_config.json b/post_config.json new file mode 100644 index 0000000000000000000000000000000000000000..8e1d6b0aa669ce6b2fb506209abc447289fd8331 --- /dev/null +++ b/post_config.json @@ -0,0 +1,14 @@ +{ + "enable_temperature" : false, + "temperature" : 0.9, + + "enable_repetition_penalty" : false, + "repetition_penalty" : 1.2, + "penalty_window" : 20, + + "enable_top_p_sampling" : false, + "top_p" : 0.8, + + "enable_top_k_sampling" : false, + "top_k" : 10 +} \ No newline at end of file diff --git a/qwen2.5_tokenizer/README.md b/qwen2.5_tokenizer/README.md new file mode 100644 index 0000000000000000000000000000000000000000..101d3b50badae6ff74d17582cbcda8a5e47b5478 --- /dev/null +++ b/qwen2.5_tokenizer/README.md @@ -0,0 +1,3 @@ +# LLM + +Standalone Qwen2.5-3B language model split from LocateAnything. Original keys stripped the `language_model.` prefix. diff --git a/qwen2.5_tokenizer/added_tokens.json b/qwen2.5_tokenizer/added_tokens.json new file mode 100644 index 0000000000000000000000000000000000000000..c80e343d9ad28395ddd247fc3960e66590413846 --- /dev/null +++ b/qwen2.5_tokenizer/added_tokens.json @@ -0,0 +1,1040 @@ +{ + "": 151669, + "": 152680, + "": 151667, + "": 151675, + "": 151671, + "": 151673, + "": 151658, + "<0>": 151677, + "<1000>": 152677, + "<100>": 151777, + "<101>": 151778, + "<102>": 151779, + "<103>": 151780, + "<104>": 151781, + "<105>": 151782, + "<106>": 151783, + "<107>": 151784, + "<108>": 151785, + "<109>": 151786, + "<10>": 151687, + "<110>": 151787, + "<111>": 151788, + "<112>": 151789, + "<113>": 151790, + "<114>": 151791, + "<115>": 151792, + "<116>": 151793, + "<117>": 151794, + "<118>": 151795, + "<119>": 151796, + "<11>": 151688, + "<120>": 151797, + "<121>": 151798, + "<122>": 151799, + "<123>": 151800, + "<124>": 151801, + "<125>": 151802, + "<126>": 151803, + "<127>": 151804, + "<128>": 151805, + "<129>": 151806, + "<12>": 151689, + "<130>": 151807, + "<131>": 151808, + "<132>": 151809, + "<133>": 151810, + "<134>": 151811, + "<135>": 151812, + "<136>": 151813, + "<137>": 151814, + "<138>": 151815, + "<139>": 151816, + "<13>": 151690, + "<140>": 151817, + "<141>": 151818, + "<142>": 151819, + "<143>": 151820, + "<144>": 151821, + "<145>": 151822, + "<146>": 151823, + "<147>": 151824, + "<148>": 151825, + "<149>": 151826, + "<14>": 151691, + "<150>": 151827, + "<151>": 151828, + "<152>": 151829, + "<153>": 151830, + "<154>": 151831, + "<155>": 151832, + "<156>": 151833, + "<157>": 151834, + "<158>": 151835, + "<159>": 151836, + "<15>": 151692, + "<160>": 151837, + "<161>": 151838, + "<162>": 151839, + "<163>": 151840, + "<164>": 151841, + "<165>": 151842, + "<166>": 151843, + "<167>": 151844, + "<168>": 151845, + "<169>": 151846, + "<16>": 151693, + "<170>": 151847, + "<171>": 151848, + "<172>": 151849, + "<173>": 151850, + "<174>": 151851, + "<175>": 151852, + "<176>": 151853, + "<177>": 151854, + "<178>": 151855, + "<179>": 151856, + "<17>": 151694, + "<180>": 151857, + "<181>": 151858, + "<182>": 151859, + "<183>": 151860, + "<184>": 151861, + "<185>": 151862, + "<186>": 151863, + "<187>": 151864, + "<188>": 151865, + "<189>": 151866, + "<18>": 151695, + "<190>": 151867, + "<191>": 151868, + "<192>": 151869, + "<193>": 151870, + "<194>": 151871, + "<195>": 151872, + "<196>": 151873, + "<197>": 151874, + "<198>": 151875, + "<199>": 151876, + "<19>": 151696, + "<1>": 151678, + "<200>": 151877, + "<201>": 151878, + "<202>": 151879, + "<203>": 151880, + "<204>": 151881, + "<205>": 151882, + "<206>": 151883, + "<207>": 151884, + "<208>": 151885, + "<209>": 151886, + "<20>": 151697, + "<210>": 151887, + "<211>": 151888, + "<212>": 151889, + "<213>": 151890, + "<214>": 151891, + "<215>": 151892, + "<216>": 151893, + "<217>": 151894, + "<218>": 151895, + "<219>": 151896, + "<21>": 151698, + "<220>": 151897, + "<221>": 151898, + "<222>": 151899, + "<223>": 151900, + "<224>": 151901, + "<225>": 151902, + "<226>": 151903, + "<227>": 151904, + "<228>": 151905, + "<229>": 151906, + "<22>": 151699, + "<230>": 151907, + "<231>": 151908, + "<232>": 151909, + "<233>": 151910, + "<234>": 151911, + "<235>": 151912, + "<236>": 151913, + "<237>": 151914, + "<238>": 151915, + "<239>": 151916, + "<23>": 151700, + "<240>": 151917, + "<241>": 151918, + "<242>": 151919, + "<243>": 151920, + "<244>": 151921, + "<245>": 151922, + "<246>": 151923, + "<247>": 151924, + "<248>": 151925, + "<249>": 151926, + "<24>": 151701, + "<250>": 151927, + "<251>": 151928, + "<252>": 151929, + "<253>": 151930, + "<254>": 151931, + "<255>": 151932, + "<256>": 151933, + "<257>": 151934, + "<258>": 151935, + "<259>": 151936, + "<25>": 151702, + "<260>": 151937, + "<261>": 151938, + "<262>": 151939, + "<263>": 151940, + "<264>": 151941, + "<265>": 151942, + "<266>": 151943, + "<267>": 151944, + "<268>": 151945, + "<269>": 151946, + "<26>": 151703, + "<270>": 151947, + "<271>": 151948, + "<272>": 151949, + "<273>": 151950, + "<274>": 151951, + "<275>": 151952, + "<276>": 151953, + "<277>": 151954, + "<278>": 151955, + "<279>": 151956, + "<27>": 151704, + "<280>": 151957, + "<281>": 151958, + "<282>": 151959, + "<283>": 151960, + "<284>": 151961, + "<285>": 151962, + "<286>": 151963, + "<287>": 151964, + "<288>": 151965, + "<289>": 151966, + "<28>": 151705, + "<290>": 151967, + "<291>": 151968, + "<292>": 151969, + "<293>": 151970, + "<294>": 151971, + "<295>": 151972, + "<296>": 151973, + "<297>": 151974, + "<298>": 151975, + "<299>": 151976, + "<29>": 151706, + "<2>": 151679, + "<300>": 151977, + "<301>": 151978, + "<302>": 151979, + "<303>": 151980, + "<304>": 151981, + "<305>": 151982, + "<306>": 151983, + "<307>": 151984, + "<308>": 151985, + "<309>": 151986, + "<30>": 151707, + "<310>": 151987, + "<311>": 151988, + "<312>": 151989, + "<313>": 151990, + "<314>": 151991, + "<315>": 151992, + "<316>": 151993, + "<317>": 151994, + "<318>": 151995, + "<319>": 151996, + "<31>": 151708, + "<320>": 151997, + "<321>": 151998, + "<322>": 151999, + "<323>": 152000, + "<324>": 152001, + "<325>": 152002, + "<326>": 152003, + "<327>": 152004, + "<328>": 152005, + "<329>": 152006, + "<32>": 151709, + "<330>": 152007, + "<331>": 152008, + "<332>": 152009, + "<333>": 152010, + "<334>": 152011, + "<335>": 152012, + "<336>": 152013, + "<337>": 152014, + "<338>": 152015, + "<339>": 152016, + "<33>": 151710, + "<340>": 152017, + "<341>": 152018, + "<342>": 152019, + "<343>": 152020, + "<344>": 152021, + "<345>": 152022, + "<346>": 152023, + "<347>": 152024, + "<348>": 152025, + "<349>": 152026, + "<34>": 151711, + "<350>": 152027, + "<351>": 152028, + "<352>": 152029, + "<353>": 152030, + "<354>": 152031, + "<355>": 152032, + "<356>": 152033, + "<357>": 152034, + "<358>": 152035, + "<359>": 152036, + "<35>": 151712, + "<360>": 152037, + "<361>": 152038, + "<362>": 152039, + "<363>": 152040, + "<364>": 152041, + "<365>": 152042, + "<366>": 152043, + "<367>": 152044, + "<368>": 152045, + "<369>": 152046, + "<36>": 151713, + "<370>": 152047, + "<371>": 152048, + "<372>": 152049, + "<373>": 152050, + "<374>": 152051, + "<375>": 152052, + "<376>": 152053, + "<377>": 152054, + "<378>": 152055, + "<379>": 152056, + "<37>": 151714, + "<380>": 152057, + "<381>": 152058, + "<382>": 152059, + "<383>": 152060, + "<384>": 152061, + "<385>": 152062, + "<386>": 152063, + "<387>": 152064, + "<388>": 152065, + "<389>": 152066, + "<38>": 151715, + "<390>": 152067, + "<391>": 152068, + "<392>": 152069, + "<393>": 152070, + "<394>": 152071, + "<395>": 152072, + "<396>": 152073, + "<397>": 152074, + "<398>": 152075, + "<399>": 152076, + "<39>": 151716, + "<3>": 151680, + "<400>": 152077, + "<401>": 152078, + "<402>": 152079, + "<403>": 152080, + "<404>": 152081, + "<405>": 152082, + "<406>": 152083, + "<407>": 152084, + "<408>": 152085, + "<409>": 152086, + "<40>": 151717, + "<410>": 152087, + "<411>": 152088, + "<412>": 152089, + "<413>": 152090, + "<414>": 152091, + "<415>": 152092, + "<416>": 152093, + "<417>": 152094, + "<418>": 152095, + "<419>": 152096, + "<41>": 151718, + "<420>": 152097, + "<421>": 152098, + "<422>": 152099, + "<423>": 152100, + "<424>": 152101, + "<425>": 152102, + "<426>": 152103, + "<427>": 152104, + "<428>": 152105, + "<429>": 152106, + "<42>": 151719, + "<430>": 152107, + "<431>": 152108, + "<432>": 152109, + "<433>": 152110, + "<434>": 152111, + "<435>": 152112, + "<436>": 152113, + "<437>": 152114, + "<438>": 152115, + "<439>": 152116, + "<43>": 151720, + "<440>": 152117, + "<441>": 152118, + "<442>": 152119, + "<443>": 152120, + "<444>": 152121, + "<445>": 152122, + "<446>": 152123, + "<447>": 152124, + "<448>": 152125, + "<449>": 152126, + "<44>": 151721, + "<450>": 152127, + "<451>": 152128, + "<452>": 152129, + "<453>": 152130, + "<454>": 152131, + "<455>": 152132, + "<456>": 152133, + "<457>": 152134, + "<458>": 152135, + "<459>": 152136, + "<45>": 151722, + "<460>": 152137, + "<461>": 152138, + "<462>": 152139, + "<463>": 152140, + "<464>": 152141, + "<465>": 152142, + "<466>": 152143, + "<467>": 152144, + "<468>": 152145, + "<469>": 152146, + "<46>": 151723, + "<470>": 152147, + "<471>": 152148, + "<472>": 152149, + "<473>": 152150, + "<474>": 152151, + "<475>": 152152, + "<476>": 152153, + "<477>": 152154, + "<478>": 152155, + "<479>": 152156, + "<47>": 151724, + "<480>": 152157, + "<481>": 152158, + "<482>": 152159, + "<483>": 152160, + "<484>": 152161, + "<485>": 152162, + "<486>": 152163, + "<487>": 152164, + "<488>": 152165, + "<489>": 152166, + "<48>": 151725, + "<490>": 152167, + "<491>": 152168, + "<492>": 152169, + "<493>": 152170, + "<494>": 152171, + "<495>": 152172, + "<496>": 152173, + "<497>": 152174, + "<498>": 152175, + "<499>": 152176, + "<49>": 151726, + "<4>": 151681, + "<500>": 152177, + "<501>": 152178, + "<502>": 152179, + "<503>": 152180, + "<504>": 152181, + "<505>": 152182, + "<506>": 152183, + "<507>": 152184, + "<508>": 152185, + "<509>": 152186, + "<50>": 151727, + "<510>": 152187, + "<511>": 152188, + "<512>": 152189, + "<513>": 152190, + "<514>": 152191, + "<515>": 152192, + "<516>": 152193, + "<517>": 152194, + "<518>": 152195, + "<519>": 152196, + "<51>": 151728, + "<520>": 152197, + "<521>": 152198, + "<522>": 152199, + "<523>": 152200, + "<524>": 152201, + "<525>": 152202, + "<526>": 152203, + "<527>": 152204, + "<528>": 152205, + "<529>": 152206, + "<52>": 151729, + "<530>": 152207, + "<531>": 152208, + "<532>": 152209, + "<533>": 152210, + "<534>": 152211, + "<535>": 152212, + "<536>": 152213, + "<537>": 152214, + "<538>": 152215, + "<539>": 152216, + "<53>": 151730, + "<540>": 152217, + "<541>": 152218, + "<542>": 152219, + "<543>": 152220, + "<544>": 152221, + "<545>": 152222, + "<546>": 152223, + "<547>": 152224, + "<548>": 152225, + "<549>": 152226, + "<54>": 151731, + "<550>": 152227, + "<551>": 152228, + "<552>": 152229, + "<553>": 152230, + "<554>": 152231, + "<555>": 152232, + "<556>": 152233, + "<557>": 152234, + "<558>": 152235, + "<559>": 152236, + "<55>": 151732, + "<560>": 152237, + "<561>": 152238, + "<562>": 152239, + "<563>": 152240, + "<564>": 152241, + "<565>": 152242, + "<566>": 152243, + "<567>": 152244, + "<568>": 152245, + "<569>": 152246, + "<56>": 151733, + "<570>": 152247, + "<571>": 152248, + "<572>": 152249, + "<573>": 152250, + "<574>": 152251, + "<575>": 152252, + "<576>": 152253, + "<577>": 152254, + "<578>": 152255, + "<579>": 152256, + "<57>": 151734, + "<580>": 152257, + "<581>": 152258, + "<582>": 152259, + "<583>": 152260, + "<584>": 152261, + "<585>": 152262, + "<586>": 152263, + "<587>": 152264, + "<588>": 152265, + "<589>": 152266, + "<58>": 151735, + "<590>": 152267, + "<591>": 152268, + "<592>": 152269, + "<593>": 152270, + "<594>": 152271, + "<595>": 152272, + "<596>": 152273, + "<597>": 152274, + "<598>": 152275, + "<599>": 152276, + "<59>": 151736, + "<5>": 151682, + "<600>": 152277, + "<601>": 152278, + "<602>": 152279, + "<603>": 152280, + "<604>": 152281, + "<605>": 152282, + "<606>": 152283, + "<607>": 152284, + "<608>": 152285, + "<609>": 152286, + "<60>": 151737, + "<610>": 152287, + "<611>": 152288, + "<612>": 152289, + "<613>": 152290, + "<614>": 152291, + "<615>": 152292, + "<616>": 152293, + "<617>": 152294, + "<618>": 152295, + "<619>": 152296, + "<61>": 151738, + "<620>": 152297, + "<621>": 152298, + "<622>": 152299, + "<623>": 152300, + "<624>": 152301, + "<625>": 152302, + "<626>": 152303, + "<627>": 152304, + "<628>": 152305, + "<629>": 152306, + "<62>": 151739, + "<630>": 152307, + "<631>": 152308, + "<632>": 152309, + "<633>": 152310, + "<634>": 152311, + "<635>": 152312, + "<636>": 152313, + "<637>": 152314, + "<638>": 152315, + "<639>": 152316, + "<63>": 151740, + "<640>": 152317, + "<641>": 152318, + "<642>": 152319, + "<643>": 152320, + "<644>": 152321, + "<645>": 152322, + "<646>": 152323, + "<647>": 152324, + "<648>": 152325, + "<649>": 152326, + "<64>": 151741, + "<650>": 152327, + "<651>": 152328, + "<652>": 152329, + "<653>": 152330, + "<654>": 152331, + "<655>": 152332, + "<656>": 152333, + "<657>": 152334, + "<658>": 152335, + "<659>": 152336, + "<65>": 151742, + "<660>": 152337, + "<661>": 152338, + "<662>": 152339, + "<663>": 152340, + "<664>": 152341, + "<665>": 152342, + "<666>": 152343, + "<667>": 152344, + "<668>": 152345, + "<669>": 152346, + "<66>": 151743, + "<670>": 152347, + "<671>": 152348, + "<672>": 152349, + "<673>": 152350, + "<674>": 152351, + "<675>": 152352, + "<676>": 152353, + "<677>": 152354, + "<678>": 152355, + "<679>": 152356, + "<67>": 151744, + "<680>": 152357, + "<681>": 152358, + "<682>": 152359, + "<683>": 152360, + "<684>": 152361, + "<685>": 152362, + "<686>": 152363, + "<687>": 152364, + "<688>": 152365, + "<689>": 152366, + "<68>": 151745, + "<690>": 152367, + "<691>": 152368, + "<692>": 152369, + "<693>": 152370, + "<694>": 152371, + "<695>": 152372, + "<696>": 152373, + "<697>": 152374, + "<698>": 152375, + "<699>": 152376, + "<69>": 151746, + "<6>": 151683, + "<700>": 152377, + "<701>": 152378, + "<702>": 152379, + "<703>": 152380, + "<704>": 152381, + "<705>": 152382, + "<706>": 152383, + "<707>": 152384, + "<708>": 152385, + "<709>": 152386, + "<70>": 151747, + "<710>": 152387, + "<711>": 152388, + "<712>": 152389, + "<713>": 152390, + "<714>": 152391, + "<715>": 152392, + "<716>": 152393, + "<717>": 152394, + "<718>": 152395, + "<719>": 152396, + "<71>": 151748, + "<720>": 152397, + "<721>": 152398, + "<722>": 152399, + "<723>": 152400, + "<724>": 152401, + "<725>": 152402, + "<726>": 152403, + "<727>": 152404, + "<728>": 152405, + "<729>": 152406, + "<72>": 151749, + "<730>": 152407, + "<731>": 152408, + "<732>": 152409, + "<733>": 152410, + "<734>": 152411, + "<735>": 152412, + "<736>": 152413, + "<737>": 152414, + "<738>": 152415, + "<739>": 152416, + "<73>": 151750, + "<740>": 152417, + "<741>": 152418, + "<742>": 152419, + "<743>": 152420, + "<744>": 152421, + "<745>": 152422, + "<746>": 152423, + "<747>": 152424, + "<748>": 152425, + "<749>": 152426, + "<74>": 151751, + "<750>": 152427, + "<751>": 152428, + "<752>": 152429, + "<753>": 152430, + "<754>": 152431, + "<755>": 152432, + "<756>": 152433, + "<757>": 152434, + "<758>": 152435, + "<759>": 152436, + "<75>": 151752, + "<760>": 152437, + "<761>": 152438, + "<762>": 152439, + "<763>": 152440, + "<764>": 152441, + "<765>": 152442, + "<766>": 152443, + "<767>": 152444, + "<768>": 152445, + "<769>": 152446, + "<76>": 151753, + "<770>": 152447, + "<771>": 152448, + "<772>": 152449, + "<773>": 152450, + "<774>": 152451, + "<775>": 152452, + "<776>": 152453, + "<777>": 152454, + "<778>": 152455, + "<779>": 152456, + "<77>": 151754, + "<780>": 152457, + "<781>": 152458, + "<782>": 152459, + "<783>": 152460, + "<784>": 152461, + "<785>": 152462, + "<786>": 152463, + "<787>": 152464, + "<788>": 152465, + "<789>": 152466, + "<78>": 151755, + "<790>": 152467, + "<791>": 152468, + "<792>": 152469, + "<793>": 152470, + "<794>": 152471, + "<795>": 152472, + "<796>": 152473, + "<797>": 152474, + "<798>": 152475, + "<799>": 152476, + "<79>": 151756, + "<7>": 151684, + "<800>": 152477, + "<801>": 152478, + "<802>": 152479, + "<803>": 152480, + "<804>": 152481, + "<805>": 152482, + "<806>": 152483, + "<807>": 152484, + "<808>": 152485, + "<809>": 152486, + "<80>": 151757, + "<810>": 152487, + "<811>": 152488, + "<812>": 152489, + "<813>": 152490, + "<814>": 152491, + "<815>": 152492, + "<816>": 152493, + "<817>": 152494, + "<818>": 152495, + "<819>": 152496, + "<81>": 151758, + "<820>": 152497, + "<821>": 152498, + "<822>": 152499, + "<823>": 152500, + "<824>": 152501, + "<825>": 152502, + "<826>": 152503, + "<827>": 152504, + "<828>": 152505, + "<829>": 152506, + "<82>": 151759, + "<830>": 152507, + "<831>": 152508, + "<832>": 152509, + "<833>": 152510, + "<834>": 152511, + "<835>": 152512, + "<836>": 152513, + "<837>": 152514, + "<838>": 152515, + "<839>": 152516, + "<83>": 151760, + "<840>": 152517, + "<841>": 152518, + "<842>": 152519, + "<843>": 152520, + "<844>": 152521, + "<845>": 152522, + "<846>": 152523, + "<847>": 152524, + "<848>": 152525, + "<849>": 152526, + "<84>": 151761, + "<850>": 152527, + "<851>": 152528, + "<852>": 152529, + "<853>": 152530, + "<854>": 152531, + "<855>": 152532, + "<856>": 152533, + "<857>": 152534, + "<858>": 152535, + "<859>": 152536, + "<85>": 151762, + "<860>": 152537, + "<861>": 152538, + "<862>": 152539, + "<863>": 152540, + "<864>": 152541, + "<865>": 152542, + "<866>": 152543, + "<867>": 152544, + "<868>": 152545, + "<869>": 152546, + "<86>": 151763, + "<870>": 152547, + "<871>": 152548, + "<872>": 152549, + "<873>": 152550, + "<874>": 152551, + "<875>": 152552, + "<876>": 152553, + "<877>": 152554, + "<878>": 152555, + "<879>": 152556, + "<87>": 151764, + "<880>": 152557, + "<881>": 152558, + "<882>": 152559, + "<883>": 152560, + "<884>": 152561, + "<885>": 152562, + "<886>": 152563, + "<887>": 152564, + "<888>": 152565, + "<889>": 152566, + "<88>": 151765, + "<890>": 152567, + "<891>": 152568, + "<892>": 152569, + "<893>": 152570, + "<894>": 152571, + "<895>": 152572, + "<896>": 152573, + "<897>": 152574, + "<898>": 152575, + "<899>": 152576, + "<89>": 151766, + "<8>": 151685, + "<900>": 152577, + "<901>": 152578, + "<902>": 152579, + "<903>": 152580, + "<904>": 152581, + "<905>": 152582, + "<906>": 152583, + "<907>": 152584, + "<908>": 152585, + "<909>": 152586, + "<90>": 151767, + "<910>": 152587, + "<911>": 152588, + "<912>": 152589, + "<913>": 152590, + "<914>": 152591, + "<915>": 152592, + "<916>": 152593, + "<917>": 152594, + "<918>": 152595, + "<919>": 152596, + "<91>": 151768, + "<920>": 152597, + "<921>": 152598, + "<922>": 152599, + "<923>": 152600, + "<924>": 152601, + "<925>": 152602, + "<926>": 152603, + "<927>": 152604, + "<928>": 152605, + "<929>": 152606, + "<92>": 151769, + "<930>": 152607, + "<931>": 152608, + "<932>": 152609, + "<933>": 152610, + "<934>": 152611, + "<935>": 152612, + "<936>": 152613, + "<937>": 152614, + "<938>": 152615, + "<939>": 152616, + "<93>": 151770, + "<940>": 152617, + "<941>": 152618, + "<942>": 152619, + "<943>": 152620, + "<944>": 152621, + "<945>": 152622, + "<946>": 152623, + "<947>": 152624, + "<948>": 152625, + "<949>": 152626, + "<94>": 151771, + "<950>": 152627, + "<951>": 152628, + "<952>": 152629, + "<953>": 152630, + "<954>": 152631, + "<955>": 152632, + "<956>": 152633, + "<957>": 152634, + "<958>": 152635, + "<959>": 152636, + "<95>": 151772, + "<960>": 152637, + "<961>": 152638, + "<962>": 152639, + "<963>": 152640, + "<964>": 152641, + "<965>": 152642, + "<966>": 152643, + "<967>": 152644, + "<968>": 152645, + "<969>": 152646, + "<96>": 151773, + "<970>": 152647, + "<971>": 152648, + "<972>": 152649, + "<973>": 152650, + "<974>": 152651, + "<975>": 152652, + "<976>": 152653, + "<977>": 152654, + "<978>": 152655, + "<979>": 152656, + "<97>": 151774, + "<980>": 152657, + "<981>": 152658, + "<982>": 152659, + "<983>": 152660, + "<984>": 152661, + "<985>": 152662, + "<986>": 152663, + "<987>": 152664, + "<988>": 152665, + "<989>": 152666, + "<98>": 151775, + "<990>": 152667, + "<991>": 152668, + "<992>": 152669, + "<993>": 152670, + "<994>": 152671, + "<995>": 152672, + "<996>": 152673, + "<997>": 152674, + "<998>": 152675, + "<999>": 152676, + "<99>": 151776, + "<9>": 151686, + "": 151665, + "": 151668, + "": 151666, + "": 151674, + "": 152678, + "": 151670, + "": 151672, + "": 152679, + "": 151676, + "": 151657, + "<|box_end|>": 151649, + "<|box_start|>": 151648, + "<|endoftext|>": 151643, + "<|file_sep|>": 151664, + "<|fim_middle|>": 151660, + "<|fim_pad|>": 151662, + "<|fim_prefix|>": 151659, + "<|fim_suffix|>": 151661, + "<|im_end|>": 151645, + "<|im_start|>": 151644, + "<|image_pad|>": 151655, + "<|object_ref_end|>": 151647, + "<|object_ref_start|>": 151646, + "<|quad_end|>": 151651, + "<|quad_start|>": 151650, + "<|repo_name|>": 151663, + "<|video_pad|>": 151656, + "<|vision_end|>": 151653, + "<|vision_pad|>": 151654, + "<|vision_start|>": 151652 +} diff --git a/qwen2.5_tokenizer/chat_template.jinja b/qwen2.5_tokenizer/chat_template.jinja new file mode 100644 index 0000000000000000000000000000000000000000..bdf7919a96cfe43d50914a007b9c0877bd0ec27e --- /dev/null +++ b/qwen2.5_tokenizer/chat_template.jinja @@ -0,0 +1,54 @@ +{%- if tools %} + {{- '<|im_start|>system\n' }} + {%- if messages[0]['role'] == 'system' %} + {{- messages[0]['content'] }} + {%- else %} + {{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }} + {%- endif %} + {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|im_end|>\n" }} +{%- else %} + {%- if messages[0]['role'] == 'system' %} + {{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }} + {%- else %} + {{- '<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n' }} + {%- endif %} +{%- endif %} +{%- for message in messages %} + {%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %} + {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {{- '<|im_start|>' + message.role }} + {%- if message.content %} + {{- '\n' + message.content }} + {%- endif %} + {%- for tool_call in message.tool_calls %} + {%- if tool_call.function is defined %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {{- '\n\n{"name": "' }} + {{- tool_call.name }} + {{- '", "arguments": ' }} + {{- tool_call.arguments | tojson }} + {{- '}\n' }} + {%- endfor %} + {{- '<|im_end|>\n' }} + {%- elif message.role == "tool" %} + {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- message.content }} + {{- '\n' }} + {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %} + {{- '<|im_end|>\n' }} + {%- endif %} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} +{%- endif %} diff --git a/qwen2.5_tokenizer/chat_template.json b/qwen2.5_tokenizer/chat_template.json new file mode 100644 index 0000000000000000000000000000000000000000..701c073525fad2293bc68624280da29eed3f2f3d --- /dev/null +++ b/qwen2.5_tokenizer/chat_template.json @@ -0,0 +1,4 @@ +{ + "chat_template": "{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n{% endif %}<|im_start|>{{ message['role'] }}\n{% if message['content'] is string %}{{ message['content'] }}<|im_end|>\n{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}{% endif %}{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}