dots/mocr example

gistfile1.txt

# /// script
# requires-python = ">=3.12"
# dependencies = [ "httpx", "Pillow" ]
# ///
import asyncio
import base64
import os
import pathlib
import random
import time
from io import BytesIO

import httpx
from PIL import Image

token = os.environ["FCIO_AI_ACCESS_KEY"]

sources = []
PROMPT = "Extract the text content from this image."


def PILimage_to_base64(image, format="PNG"):
    buffered = BytesIO()
    image.save(buffered, format=format)
    base64_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    return f"data:image/{format.lower()};base64,{base64_str}"


for s in pathlib.Path(".").glob("*.png"):
    image = Image.open(s.name)
    request = {
        "model": "rednote-hilab:dots.mocr",
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {"url": PILimage_to_base64(image)},
                    },
                    {
                        "type": "text",
                        "text": f"<|img|><|imgpad|><|endofimg|>{PROMPT}",
                    },  # if no "<|img|><|imgpad|><|endofimg|>" here,vllm v1 will add "\n" here
                ],
            }
        ],
        "temperature": "0.1",
        "top_p": 0.9,
        "max_completion_tokens": 32768,
    }
    sources.append(request)

headers = {"Authorization": f"Bearer {token}"}

print(f"Loaded {len(sources)} requests.")

TOTAL_TOKENS = 0
SEM = asyncio.Semaphore(32)
IN_PROGRESS = 0
SUCCESS = 0
ERRORS = 0


async def run_request():
    global TOTAL_TOKENS, IN_PROGRESS, SUCCESS, ERRORS
    async with SEM:
        IN_PROGRESS += 1
        try:
            async with httpx.AsyncClient(timeout=120) as client:
                request = random.choice(sources)
                print("progress:", IN_PROGRESS)
                r = await client.post(
                    "https://ai.rzob.fcio.net/openai/v1/chat/completions",
                    json=request,
                    headers=headers,
                )
                if not r.status_code == 200:
                    print("error", r.status_code)
                    print(r.content)
                    ERRORS += 1
                    return
                x = r.json()

                if not "usage" in x:
                    print(request["filename"])
                    print(x)
                    return
                print(x["usage"]["total_tokens"])
                TOTAL_TOKENS += x["usage"]["total_tokens"]
                SUCCESS += 1
        finally:
            IN_PROGRESS -= 1


async def main():
    start = time.time()
    try:
        tasks = []
        for x in range(1000):
            tasks.append(asyncio.create_task(run_request()))
        await asyncio.gather(*tasks)
    finally:
        end = time.time()
        print("duration: ", end - start)
        print("successfull requests: ", SUCCESS)
        print("successfull requests/s: ", SUCCESS / (end - start))
        print("tokens: ", TOTAL_TOKENS)
        print("tokens/s: ", TOTAL_TOKENS / (end - start))


if __name__ == "__main__":
    asyncio.run(main())