Spaces:
Sleeping
Sleeping
| """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.""" | |
| # Common delimiters used by chat templates | |
| 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) | |