| """Shared text generation utilities.""" |
| from __future__ import annotations |
|
|
| import re |
| import logging |
| from typing import List, Dict, Any |
|
|
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| from src.config import DEFAULT_MAX_NEW_TOKENS, DEFAULT_TEMPERATURE |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| def strip_think_blocks(text: str) -> str: |
| """Remove <think>...</think> blocks if the model emits them.""" |
| return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip() |
|
|
|
|
| def extract_assistant_reply(text: str) -> str: |
| """Extract the assistant reply from a chat-templated output string.""" |
| |
| markers = [ |
| "assistant\n", |
| "assistant", |
| "<|im_start|>assistant", |
| "<|start_header_id|>assistant<|end_header_id|>", |
| ] |
| for marker in markers: |
| if marker in text: |
| reply = text.split(marker, 1)[-1] |
| return strip_think_blocks(reply).strip() |
| return strip_think_blocks(text).strip() |
|
|
|
|
| def chat_generate( |
| model: AutoModelForCausalLM, |
| tokenizer: AutoTokenizer, |
| messages: List[Dict[str, Any]], |
| max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS, |
| temperature: float = DEFAULT_TEMPERATURE, |
| do_sample: bool = False, |
| ) -> str: |
| """Generate an assistant reply from a list of chat messages.""" |
| prompt = tokenizer.apply_chat_template( |
| messages, |
| tokenize=False, |
| add_generation_prompt=True, |
| ) |
|
|
| inputs = tokenizer( |
| prompt, |
| return_tensors="pt", |
| padding=True, |
| truncation=True, |
| ) |
| if hasattr(model, "device") and model.device.type != "meta": |
| inputs = {k: v.to(model.device) for k, v in inputs.items()} |
|
|
| with torch.no_grad(): |
| outputs = model.generate( |
| **inputs, |
| max_new_tokens=max_new_tokens, |
| do_sample=do_sample, |
| temperature=temperature if do_sample else None, |
| pad_token_id=tokenizer.pad_token_id, |
| eos_token_id=tokenizer.eos_token_id, |
| ) |
|
|
| generated = outputs[0][inputs["input_ids"].shape[1]:] |
| text = tokenizer.decode(generated, skip_special_tokens=True) |
| return extract_assistant_reply(text) |
|
|
