training_coach/parser_runtime.py

from functools import lru_cache
import logging
import os
from time import perf_counter

from training_coach.parser import (
    PARSER_MODEL,
    build_parser_messages,
    log_parser_messages,
    log_parser_response_text,
    parse_model_response,
)
from training_coach.models import ParsedCheckIn


MODEL_CACHE_ENV_VAR = "PARSER_MODEL_CACHE_DIR"
DEFAULT_TRANSFORMERS_MODEL = "Qwen/Qwen3-1.7B"
DEFAULT_MAX_NEW_TOKENS = 384
DEFAULT_ZEROGPU_DURATION_SECONDS = 120
logger = logging.getLogger(__name__)


class ParserRuntimeUnavailableError(RuntimeError):
    pass


def _gpu_decorator():
    try:
        import spaces
    except ImportError:
        return lambda function: function
    duration = int(
        os.getenv("ZEROGPU_DURATION_SECONDS", str(DEFAULT_ZEROGPU_DURATION_SECONDS))
    )
    return spaces.GPU(duration=duration)


def _load_transformers():
    try:
        from transformers import AutoModelForCausalLM, AutoTokenizer
    except ImportError as error:
        raise ParserRuntimeUnavailableError(
            "Install transformers, torch, and accelerate to run the parser model."
        ) from error

    return AutoModelForCausalLM, AutoTokenizer


@lru_cache(maxsize=1)
def load_parser_model(model_name: str = DEFAULT_TRANSFORMERS_MODEL):
    start_time = perf_counter()
    logger.info("event=parser_model_load_start model=%s", model_name)
    AutoModelForCausalLM, AutoTokenizer = _load_transformers()
    cache_dir = os.getenv(MODEL_CACHE_ENV_VAR) or None
    tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        cache_dir=cache_dir,
        device_map="auto",
        torch_dtype="auto",
    )
    logger.info(
        "event=parser_model_load_complete model=%s cache_dir_configured=%s elapsed_ms=%s",
        model_name,
        cache_dir is not None,
        round((perf_counter() - start_time) * 1000),
    )
    return tokenizer, model


@_gpu_decorator()
def generate_parser_response(
    raw_text: str,
    model_name: str = DEFAULT_TRANSFORMERS_MODEL,
) -> str:
    start_time = perf_counter()
    logger.info(
        "event=parser_generate_start model=%s text_chars=%s",
        model_name,
        len(raw_text),
    )
    tokenizer, model = load_parser_model(model_name)
    messages = build_parser_messages(raw_text)
    log_parser_messages(
        backend="transformers",
        model_name=model_name,
        messages=messages,
    )
    chat_template_kwargs = {
        "tokenize": False,
        "add_generation_prompt": True,
        "enable_thinking": False,
    }
    try:
        prompt = tokenizer.apply_chat_template(messages, **chat_template_kwargs)
    except TypeError:
        chat_template_kwargs.pop("enable_thinking")
        prompt = tokenizer.apply_chat_template(messages, **chat_template_kwargs)
    inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
    input_token_count = inputs.input_ids.shape[-1]
    output_ids = model.generate(
        **inputs,
        max_new_tokens=int(os.getenv("PARSER_MAX_NEW_TOKENS", str(DEFAULT_MAX_NEW_TOKENS))),
        do_sample=False,
        pad_token_id=tokenizer.eos_token_id,
    )
    generated_ids = output_ids[:, inputs.input_ids.shape[-1] :]
    generated_token_count = generated_ids.shape[-1]
    response_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
    logger.info(
        "event=parser_generate_complete model=%s prompt_chars=%s "
        "input_tokens=%s generated_tokens=%s response_chars=%s elapsed_ms=%s",
        model_name,
        len(prompt),
        input_token_count,
        generated_token_count,
        len(response_text),
        round((perf_counter() - start_time) * 1000),
    )
    log_parser_response_text(
        backend="transformers",
        model_name=model_name,
        response_text=response_text,
    )
    return response_text


def parse_check_in_with_model(
    raw_text: str,
    model_name: str | None = None,
) -> ParsedCheckIn:
    model_name = model_name or os.getenv("PARSER_MODEL_ID", DEFAULT_TRANSFORMERS_MODEL)
    logger.info("event=parser_model_parse_start model=%s", model_name)
    response_text = generate_parser_response(raw_text, model_name=model_name)
    parsed = parse_model_response(response_text)
    logger.info(
        "event=parser_model_parse_complete model=%s missing_fields=%s follow_up_questions=%s",
        model_name,
        len(parsed.missing_fields),
        len(parsed.follow_up_questions),
    )
    return parsed