models.py

import torch
from typing import Dict, List, Optional, Tuple

from transformers import AutoModelForCausalLM, AutoTokenizer

try:
    from vllm import LLM, SamplingParams
    _HAS_VLLM = True
except ImportError:
    _HAS_VLLM = False


def _ensure_pad_token(tokenizer: AutoTokenizer) -> None:
    if tokenizer.pad_token_id is None:
        if tokenizer.eos_token is not None:
            tokenizer.pad_token = tokenizer.eos_token
        else:
            tokenizer.add_special_tokens({"pad_token": "<pad>"})


class ModelWrapper:
    """
    Thin wrapper around a HuggingFace causal-LM (with optional vLLM backend).

    Supports:
      - Text generation via HF transformers (default)
      - Text generation via vLLM (--use_vllm)
    """

    def __init__(self, model_name: str, device: torch.device, use_vllm: bool = False, args=None):
        self.model_name = model_name
        self.device = device
        self.use_vllm = use_vllm and _HAS_VLLM
        self.vllm_engine = None
        self.args = args

        if self.use_vllm:
            tp_size = max(1, int(getattr(args, "tensor_parallel_size", 1)))
            gpu_util = float(getattr(args, "gpu_memory_utilization", 0.9))
            max_model_len = getattr(args, "max_model_len", None)
            print(f"[vLLM] Loading {model_name} with tensor_parallel_size={tp_size}")
            llm_kwargs = dict(
                model=model_name,
                tensor_parallel_size=tp_size,
                gpu_memory_utilization=gpu_util,
            )
            if max_model_len is not None:
                llm_kwargs["max_model_len"] = max_model_len
            self.vllm_engine = LLM(**llm_kwargs)
            # Store actual context limit for length guards downstream
            self._max_model_len = self.vllm_engine.llm_engine.model_config.max_model_len
            self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
            _ensure_pad_token(self.tokenizer)
            return

        # ── HuggingFace path ──
        print(f"[HF] Loading {model_name} ...")
        self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
        _ensure_pad_token(self.tokenizer)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=(torch.bfloat16 if torch.cuda.is_available() else torch.float32),
        )
        if len(self.tokenizer) != self.model.get_input_embeddings().weight.shape[0]:
            self.model.resize_token_embeddings(len(self.tokenizer))
        self.model.to(device).eval()
        if hasattr(self.model.config, "use_cache"):
            self.model.config.use_cache = True

    # ── Prompt rendering ──────────────────────────────────────────────────

    def render_chat(self, messages: List[Dict], add_generation_prompt: bool = True,
                    tools: Optional[List[Dict]] = None,
                    enable_thinking: Optional[bool] = None) -> str:
        tpl = getattr(self.tokenizer, "chat_template", None)
        if tpl:
            kwargs = {"tokenize": False, "add_generation_prompt": add_generation_prompt}
            if tools is not None:
                kwargs["tools"] = tools
            if enable_thinking is not None:
                kwargs["enable_thinking"] = enable_thinking
            return self.tokenizer.apply_chat_template(messages, **kwargs)
        # Fallback for models without a chat template
        segments = []
        for msg in messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            segments.append(f"<|{role}|>\n{content}")
        if add_generation_prompt:
            segments.append("<|assistant|>")
        return "\n".join(segments)

    def prepare_chat_batch(
        self,
        batch_messages: List[List[Dict]],
        add_generation_prompt: bool = True,
        tools: Optional[List[Dict]] = None,
    ) -> Tuple[List[str], torch.Tensor, torch.Tensor, List[List[str]]]:
        prompts = [
            self.render_chat(msgs, add_generation_prompt=add_generation_prompt, tools=tools)
            for msgs in batch_messages
        ]
        encoded = self.tokenizer(
            prompts,
            return_tensors="pt",
            padding=True,
            add_special_tokens=False,
        )
        input_ids = encoded["input_ids"].to(self.device)
        attention_mask = encoded["attention_mask"].to(self.device)
        tokens_batch: List[List[str]] = []
        for i in range(len(prompts)):
            mask = attention_mask[i].bool()
            active_ids = input_ids[i][mask].tolist()
            tokens_batch.append(self.tokenizer.convert_ids_to_tokens(active_ids))
        return prompts, input_ids, attention_mask, tokens_batch

    # ── Text generation (HF) ──────────────────────────────────────────────

    @torch.no_grad()
    def generate_text_batch(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        max_new_tokens: int = 512,
        temperature: float = 0.6,
        top_p: float = 0.95,
    ) -> List[str]:
        outputs = self.model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            max_new_tokens=max_new_tokens,
            do_sample=(temperature > 0),
            temperature=temperature if temperature > 0 else 1.0,
            top_p=top_p,
            pad_token_id=self.tokenizer.pad_token_id,
            eos_token_id=self.tokenizer.eos_token_id,
        )
        prompt_len = input_ids.shape[1]
        generated: List[str] = []
        for out in outputs:
            new_tokens = out[prompt_len:]
            text = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
            generated.append(text)
        return generated

    # ── Token counting ───────────────────────────────────────────────────

    def count_tokens(self, text: str) -> int:
        """Return the number of tokens in `text` using the loaded tokenizer."""
        return len(self.tokenizer.encode(text, add_special_tokens=False))

    # ── Text generation (vLLM) ───────────────────────────────────────────

    def vllm_generate_text_batch(
        self,
        prompts: List[str],
        max_new_tokens: int = 512,
        temperature: float = 0.6,
        top_p: float = 0.95,
    ) -> List[str]:
        sampling_params = SamplingParams(
            max_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
        )
        results = self.vllm_engine.generate(prompts, sampling_params, use_tqdm=False)
        return [r.outputs[0].text for r in results]