utils.py

import base64
import json
import logging
import time
from pathlib import Path
from typing import Any, Dict, Iterable, Optional, Tuple
import re

from mimetypes import guess_type

import openai
import torch
from PIL import Image
from transformers import (
    AutoModelForCausalLM,
    AutoModelForImageTextToText,
    AutoModelForVision2Seq,
    AutoProcessor,
    AutoTokenizer,
)

def parse_batch_prompts(llm_output):
    pattern = r"<prompt>(.*?)</prompt>"
    prompts = re.findall(pattern, llm_output, re.DOTALL)
    prompts = [p.strip() for p in prompts]
    return prompts

def local_image_to_data_url(image_path: str) -> str:
    mime_type, _ = guess_type(image_path)
    if mime_type is None:
        mime_type = 'application/octet-stream'
        
    with open(image_path, "rb") as image_file:
        base64_encoded_data = base64.b64encode(image_file.read()).decode('utf-8')
    return f"data:{mime_type};base64,{base64_encoded_data}"


def get_logger(name: str = "mllm_safe_evo") -> logging.Logger:
    """
    Lightweight logger setup that is easy to import across modules.
    """
    logger = logging.getLogger(name)
    if not logger.handlers:
        handler = logging.StreamHandler()
        fmt = "%(asctime)s | %(levelname)s | %(name)s | %(message)s"
        handler.setFormatter(logging.Formatter(fmt))
        logger.addHandler(handler)
        logger.setLevel(logging.INFO)
    return logger


logger = get_logger()

# simple local model cache to avoid reloading between calls
_LOCAL_MODEL_CACHE: Dict[str, Tuple[Any, Any]] = {}
_LOCAL_VLM_CACHE: Dict[str, Tuple[Any, Any]] = {}


def _load_local_model(model: str) -> Tuple[Any, Any]:
    if torch is None or AutoTokenizer is None or AutoModelForCausalLM is None:
        raise RuntimeError("transformers/torch not available for local inference.")
    if model in _LOCAL_MODEL_CACHE:
        return _LOCAL_MODEL_CACHE[model]
    device = "cuda" if torch.cuda.is_available() else ("mps" if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available() else "cpu")
    tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    model_obj = AutoModelForCausalLM.from_pretrained(
        model,
        device_map="auto" if device in ("cuda", "mps") else None,
        trust_remote_code=True,
    )
    if device == "cpu":
        model_obj = model_obj.to(device)
    _LOCAL_MODEL_CACHE[model] = (tokenizer, model_obj)
    return tokenizer, model_obj


def _load_local_vlm(model: str) -> Tuple[Any, Any]:
    if torch is None or AutoProcessor is None or (AutoModelForImageTextToText is None and AutoModelForVision2Seq is None):
        raise RuntimeError("transformers/torch not available for local VLM inference.")
    if model in _LOCAL_VLM_CACHE:
        return _LOCAL_VLM_CACHE[model]
    device = "cuda" if torch.cuda.is_available() else ("mps" if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available() else "cpu")
    processor = AutoProcessor.from_pretrained(model, trust_remote_code=True)
    dtype = torch.float16 if torch.cuda.is_available() or device != "cpu" else None
    load_kwargs: Dict[str, Any] = {
        "trust_remote_code": True,
    }
    if device in ("cuda", "mps"):
        load_kwargs["device_map"] = "auto"
    if dtype is not None:
        load_kwargs["dtype"] = dtype
    model_obj = None
    if AutoModelForImageTextToText is not None:
        try:
            model_obj = AutoModelForImageTextToText.from_pretrained(model, **load_kwargs)
        except Exception:
            model_obj = None
    if model_obj is None and AutoModelForVision2Seq is not None:
        model_obj = AutoModelForVision2Seq.from_pretrained(model, **load_kwargs)
    if model_obj is None:
        raise RuntimeError("Failed to load local VLM model.")
    if device == "cpu":
        model_obj = model_obj.to(device)
    _LOCAL_VLM_CACHE[model] = (processor, model_obj)
    return processor, model_obj


def _make_client(model: str, api_key: str):
    if openai is None:
        raise RuntimeError("openai package not available")
    if 'gpt' in model:
        return openai.OpenAI(api_key=api_key or None)
    if 'qwen' in model or 'deepseek' in model:
        return openai.OpenAI(
            api_key=api_key or None,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )
    raise ValueError(f"Unknown model: {model}")


def call_llm(
    model: str,
    system_prompt: str,
    user_prompt: str,
    api_key: str = "",
    image_path: str = None,
    max_tokens: int = 3000,
    temperature: float = 0.1,
    enable_thinking: Optional[bool] = None,
) -> str:
    """
    Minimal chat wrapper; returns empty string if openai SDK unavailable.
    """
    if not api_key:
        # Local inference path. If image provided, try VLM; otherwise text-only.
        try:
            if image_path:
                if Image is None:
                    raise RuntimeError("PIL not available for image loading.")
                processor, model_obj = _load_local_vlm(model)
                image = Image.open(image_path).convert("RGB")
                prompt = f"{system_prompt}\n{user_prompt}" if system_prompt else user_prompt
                inputs = processor(images=image, text=prompt, return_tensors="pt").to(model_obj.device)
                with torch.no_grad():
                    output = model_obj.generate(
                        **inputs,
                        max_new_tokens=max_tokens,
                        temperature=temperature,
                    )
                return processor.decode(output[0], skip_special_tokens=True)
            else:
                tokenizer, model_obj = _load_local_model(model)
                prompt = f"{system_prompt}\n{user_prompt}" if system_prompt else user_prompt
                inputs = tokenizer(prompt, return_tensors="pt").to(model_obj.device)
                with torch.no_grad():
                    output = model_obj.generate(
                        **inputs,
                        max_new_tokens=max_tokens,
                        temperature=temperature,
                    )
                return tokenizer.decode(output[0], skip_special_tokens=True)
        except Exception as exc:  # pragma: no cover - best effort
            logger.warning("Local model inference failed (%s); returning empty response.", exc)
            return ""

    if openai is None:
        logger.warning("openai client not available; returning empty response.")
        return ""

    client = _make_client(model, api_key)
    request_kwargs: Dict[str, Any] = {}
    if enable_thinking is not None:
        request_kwargs["extra_body"] = {"enable_thinking": bool(enable_thinking),"thinking_token_limit": 2048}
    
    count = 0
    while True:
        try:
            content = [{"type": "text", "text": system_prompt}, {"type": "text", "text": user_prompt}]
            if image_path:
                content.append({
                    "type": "image_url",
                    "image_url": {"url": local_image_to_data_url(image_path)}
                })
            response = client.chat.completions.create(
                model=model,
                messages=[{"role": "user", "content": content}],
                max_tokens=max_tokens,
                temperature=temperature,
                **request_kwargs,
            )
            # 兼容新版 refusals
            if hasattr(response.choices[0].message, "refusal") and response.choices[0].message.refusal:
                return response.choices[0].message.refusal
            return response.choices[0].message.content
        except Exception as e:
            if "Error code: 400" in str(e):
                return "I'm sorry."
            else:
                count += 1
                if count == 3:
                    raise e
                time.sleep(5)
                continue
    



def save_jsonl(items: Iterable[Dict[str, Any]], path: str) -> None:
    """
    Save an iterable of dicts to JSONL.
    """
    target = Path(path)
    target.parent.mkdir(parents=True, exist_ok=True)
    with target.open("w", encoding="utf-8") as f:
        for item in items:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")


def save_model_checkpoint(model: Any, path: str) -> None:
    """
    Minimal checkpoint helper:
      - if model has .save_pretrained, call it
      - otherwise persist a JSON metadata file
    """
    target = Path(path)
    target.mkdir(parents=True, exist_ok=True)

    if hasattr(model, "save_pretrained"):
        model.save_pretrained(str(target))
        return

    metadata: Dict[str, Any] = {}
    for attr in ("model_name", "sft_steps"):
        if hasattr(model, attr):
            metadata[attr] = getattr(model, attr)

    with (target / "metadata.json").open("w", encoding="utf-8") as f:
        json.dump(metadata, f, ensure_ascii=False, indent=2)