core/gpu_generate.py

"""GPU generation function — called from app.py's @spaces.GPU wrapper."""

import re
import logging
import config

logger = logging.getLogger(__name__)

_model_cache = {}


def _load_and_generate(prompt: str, max_tokens: int = None) -> str:
    """Load model and generate text. Runs inside GPU context."""
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    max_tokens = max_tokens or config.TEXT_MODEL_MAX_TOKENS

    if "model" not in _model_cache:
        repos = [config.TEXT_MODEL_REPO, config.LIGHT_MODEL_REPO]
        for repo in repos:
            try:
                logger.info(f"Loading: {repo}")
                tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
                model = AutoModelForCausalLM.from_pretrained(
                    repo,
                    torch_dtype=torch.float16,
                    trust_remote_code=True,
                    low_cpu_mem_usage=True,
                    device_map="auto",
                )
                _model_cache["model"] = model
                _model_cache["tokenizer"] = tokenizer
                _model_cache["repo"] = repo
                logger.info(f"Loaded: {repo}")
                break
            except Exception as e:
                logger.warning(f"Failed {repo}: {e}")
                continue
        else:
            return ""

    model = _model_cache["model"]
    tokenizer = _model_cache["tokenizer"]

    system = "Output directly. No thinking. No tags. Chinese only."
    messages = [
        {"role": "system", "content": system},
        {"role": "user", "content": prompt},
    ]

    try:
        text = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
    except Exception:
        text = f"{system}\n\n{prompt}\n\nAnswer:"

    inputs = tokenizer(text, return_tensors="pt")
    inputs.pop("token_type_ids", None)

    if torch.cuda.is_available():
        inputs = {k: v.to("cuda") for k, v in inputs.items()}

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.1,
        )

    raw = tokenizer.decode(
        outputs[0][inputs["input_ids"].shape[1]:],
        skip_special_tokens=True,
    )

    thinking = re.findall(r'<think>(.*?)</think>', raw, re.DOTALL)
    raw = re.sub(r'<think>.*?</think>\s*', '', raw, flags=re.DOTALL)
    result = raw.strip()

    if not result and thinking:
        last = thinking[-1].strip()
        lines = [l.strip() for l in last.split('\n') if l.strip()]
        skip = ['嗯', '想到', '考虑', '需要', '首先', '然后', '用户']
        good = [l for l in lines if not any(l.startswith(w) for w in skip)]
        result = '\n'.join(good[-3:]) if good else '\n'.join(lines[-2:])

    logger.info(f"Generated {len(result)} chars")
    return result


# Fallback: direct call (used when app.py hasn't registered the GPU version)
def gpu_generate(prompt: str, max_tokens: int = None) -> str:
    return _load_and_generate(prompt, max_tokens)