boccaccioAI / src /inference /generate.py

Upload folder using huggingface_hub

187a491 verified 4 months ago

10.5 kB

"""BoccaccioAI inference script -- free generation and Q&A modes."""

from __future__ import annotations

import argparse
from pathlib import Path

import torch
import torch.nn.functional as F
from tokenizers import Tokenizer

from src.model.config import BoccaccioConfig
from src.model.transformer import BoccaccioForCausalLM


# ------------------------------------------------------------------ #
#  Core generation function                                          #
# ------------------------------------------------------------------ #

@torch.no_grad()
def generate(
    model: BoccaccioForCausalLM,
    input_ids: torch.Tensor,
    max_new_tokens: int = 256,
    temperature: float = 0.7,
    top_k: int = 50,
    top_p: float = 0.9,
    repetition_penalty: float = 1.0,
    eos_token_id: int | None = None,
) -> torch.Tensor:
    """Autoregressive token generation with temperature, top-k, top-p and repetition penalty.
    Parameters
    ----------
    model : BoccaccioForCausalLM
        The causal language model.
    input_ids : torch.Tensor
        Prompt token ids, shape ``(1, seq_len)``.
    max_new_tokens : int
        Maximum number of tokens to generate.
    temperature : float
        Softmax temperature (lower = more deterministic).
    top_k : int
        Keep only the *top_k* highest-probability tokens.
    top_p : float
        Nucleus sampling threshold (cumulative probability).
    repetition_penalty : float
        Penalty for tokens already generated (>1.0 reduces repetition, 1.0 = off).
    eos_token_id : int | None
        If provided, stop generation when this token is sampled.
    Returns
    -------
    torch.Tensor
        Full sequence (prompt + generated tokens), shape ``(1, seq_len + generated)``.
    """
    generated = input_ids

    for _ in range(max_new_tokens):
        # Forward pass -- we only need the logits at the last position.
        outputs = model(generated)
        logits = outputs["logits"][:, -1, :]  # (1, vocab_size)

        # Repetition penalty: reduce logits for tokens already in the sequence.
        if repetition_penalty != 1.0:
            prev_tokens = generated[0].unique()
            penalty_logits = logits[:, prev_tokens]
            # For positive logits, divide by penalty; for negative, multiply.
            penalty_logits = torch.where(
                penalty_logits > 0,
                penalty_logits / repetition_penalty,
                penalty_logits * repetition_penalty,
            )
            logits[:, prev_tokens] = penalty_logits

        # Temperature scaling.
        logits = logits / temperature

        # Top-k filtering.
        if top_k > 0:
            top_k_values, _ = torch.topk(logits, top_k, dim=-1)
            threshold = top_k_values[:, -1].unsqueeze(-1)
            logits = logits.masked_fill(logits < threshold, float("-inf"))

        # Top-p (nucleus) filtering.
        if 0.0 < top_p < 1.0:
            sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

            # Mask tokens whose cumulative probability exceeds top_p.
            sorted_mask = cumulative_probs - F.softmax(sorted_logits, dim=-1) >= top_p
            sorted_logits = sorted_logits.masked_fill(sorted_mask, float("-inf"))

            # Scatter back to original ordering.
            logits = sorted_logits.scatter(1, sorted_indices, sorted_logits)

        # Sample from the filtered distribution.
        probs = F.softmax(logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)  # (1, 1)

        generated = torch.cat([generated, next_token], dim=-1)

        if eos_token_id is not None and next_token.item() == eos_token_id:
            break

    return generated


# ------------------------------------------------------------------ #
#  CLI                                                               #
# ------------------------------------------------------------------ #

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="BoccaccioAI text generation (free generation / Q&A)",
    )

    parser.add_argument(
        "--model-dir",
        type=str,
        required=True,
        help="Path to the saved model checkpoint directory.",
    )
    parser.add_argument(
        "--tokenizer-path",
        type=str,
        default="tokenizer/boccaccio-32k.json",
        help="Path to the tokenizer JSON file.",
    )
    parser.add_argument(
        "--config-path",
        type=str,
        default="configs/model.yaml",
        help="Path to the model config YAML file.",
    )
    parser.add_argument(
        "--config-variant",
        type=str,
        default="model",
        choices=["model", "nano"],
        help="Config variant to load ('model' or 'nano').",
    )
    parser.add_argument(
        "--mode",
        type=str,
        default="generate",
        choices=["generate", "qa", "chat"],
        help="Inference mode: 'generate' for free generation, 'qa' for Q&A, 'chat' for chat.",
    )
    parser.add_argument(
        "--prompt",
        type=str,
        default=None,
        help="Text prompt for generation/chat mode.",
    )
    parser.add_argument(
        "--system-prompt",
        type=str,
        default=None,
        help="Optional system prompt for chat mode.",
    )
    parser.add_argument(
        "--context",
        type=str,
        default=None,
        help="Context text for Q&A mode (legacy).",
    )
    parser.add_argument(
        "--question",
        type=str,
        default=None,
        help="Question text for Q&A mode (legacy).",
    )
    parser.add_argument(
        "--max-new-tokens",
        type=int,
        default=256,
        help="Maximum number of new tokens to generate.",
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.7,
        help="Sampling temperature.",
    )
    parser.add_argument(
        "--top-k",
        type=int,
        default=50,
        help="Top-k filtering value.",
    )
    parser.add_argument(
        "--top-p",
        type=float,
        default=0.9,
        help="Top-p (nucleus) sampling threshold.",
    )
    parser.add_argument(
        "--repetition-penalty",
        type=float,
        default=1.2,
        help="Repetition penalty (>1.0 reduces repetition, 1.0 = off).",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="auto",
        choices=["cuda", "cpu", "auto"],
        help="Device to run inference on.",
    )

    return parser.parse_args()


# ------------------------------------------------------------------ #
#  Main                                                              #
# ------------------------------------------------------------------ #

def main() -> None:
    args = parse_args()

    # ---- device -------------------------------------------------- #
    if args.device == "auto":
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        device = torch.device(args.device)

    # ---- config -------------------------------------------------- #
    config = BoccaccioConfig.from_yaml(args.config_path, variant=args.config_variant)

    # ---- model --------------------------------------------------- #
    model = BoccaccioForCausalLM(config)

    checkpoint_dir = Path(args.model_dir)
    checkpoint_path = checkpoint_dir / "model.pt"
    state_dict = torch.load(checkpoint_path, map_location=device, weights_only=True)

    # Remap torch.compile prefix: _orig_mod.X -> X
    # Checkpoints saved with torch.compile wrap the model as _orig_mod.<inner>.
    cleaned = {}
    for key, value in state_dict.items():
        new_key = key
        if new_key.startswith("_orig_mod."):
            new_key = new_key[len("_orig_mod."):]
        cleaned[new_key] = value
    model.load_state_dict(cleaned)

    model.to(device)
    model.eval()

    # ---- tokenizer ----------------------------------------------- #
    tokenizer = Tokenizer.from_file(args.tokenizer_path)

    # ---- build prompt -------------------------------------------- #
    if args.mode == "chat":
        # Chat mode: use the same template as training (InstructionDataset).
        # Format: ### System:\n{system}\n### Utente:\n{prompt}\n### Assistente:\n
        if args.prompt is None:
            raise ValueError("Chat mode requires --prompt.")
        prompt = ""
        if args.system_prompt:
            prompt += f"### System:\n{args.system_prompt}\n"
        prompt += f"### Utente:\n{args.prompt}\n"
        prompt += f"### Assistente:\n"

    elif args.mode == "qa":
        # Legacy Q&A mode: map to chat format with context as system prompt.
        if args.context is None or args.question is None:
            raise ValueError("Q&A mode requires both --context and --question.")
        prompt = f"### System:\n{args.context}\n"
        prompt += f"### Utente:\n{args.question}\n"
        prompt += f"### Assistente:\n"

    else:
        # Free generation: raw prompt, no template.
        if args.prompt is None:
            raise ValueError("Generation mode requires --prompt.")
        prompt = args.prompt

    # ---- tokenize ------------------------------------------------ #
    encoding = tokenizer.encode(prompt)
    input_ids = torch.tensor([encoding.ids], dtype=torch.long, device=device)

    # ---- generate ------------------------------------------------ #
    eos_token_id = tokenizer.token_to_id("</s>")

    output_ids = generate(
        model,
        input_ids,
        max_new_tokens=args.max_new_tokens,
        temperature=args.temperature,
        top_k=args.top_k,
        top_p=args.top_p,
        repetition_penalty=args.repetition_penalty,
        eos_token_id=eos_token_id,
    )

    # ---- decode & print ------------------------------------------ #
    all_ids = output_ids.squeeze(0).tolist()
    prompt_len = len(encoding.ids)
    generated_ids = all_ids[prompt_len:]

    if args.mode in ("chat", "qa"):
        # In chat/qa mode, print only the generated response (not the prompt).
        response_text = tokenizer.decode(generated_ids).strip()
        print(response_text)
    else:
        # In generate mode, print the full text (prompt + generation).
        output_text = tokenizer.decode(all_ids)
        print(output_text)


if __name__ == "__main__":
    main()

	"""BoccaccioAI inference script -- free generation and Q&A modes."""

	from __future__ import annotations

	import argparse
	from pathlib import Path

	import torch
	import torch.nn.functional as F
	from tokenizers import Tokenizer

	from src.model.config import BoccaccioConfig
	from src.model.transformer import BoccaccioForCausalLM


	# ------------------------------------------------------------------ #
	# Core generation function #
	# ------------------------------------------------------------------ #

	@torch.no_grad()
	def generate(
	model: BoccaccioForCausalLM,
	input_ids: torch.Tensor,
	max_new_tokens: int = 256,
	temperature: float = 0.7,
	top_k: int = 50,
	top_p: float = 0.9,
	repetition_penalty: float = 1.0,
	eos_token_id: int \| None = None,
	) -> torch.Tensor:
	"""Autoregressive token generation with temperature, top-k, top-p and repetition penalty.

	Parameters
	----------
	model : BoccaccioForCausalLM
	The causal language model.
	input_ids : torch.Tensor
	Prompt token ids, shape ``(1, seq_len)``.
	max_new_tokens : int
	Maximum number of tokens to generate.
	temperature : float
	Softmax temperature (lower = more deterministic).
	top_k : int
	Keep only the top_k highest-probability tokens.
	top_p : float
	Nucleus sampling threshold (cumulative probability).
	repetition_penalty : float
	Penalty for tokens already generated (>1.0 reduces repetition, 1.0 = off).
	eos_token_id : int \| None
	If provided, stop generation when this token is sampled.

	Returns
	-------
	torch.Tensor
	Full sequence (prompt + generated tokens), shape ``(1, seq_len + generated)``.
	"""
	generated = input_ids

	for _ in range(max_new_tokens):
	# Forward pass -- we only need the logits at the last position.
	outputs = model(generated)
	logits = outputs["logits"][:, -1, :] # (1, vocab_size)

	# Repetition penalty: reduce logits for tokens already in the sequence.
	if repetition_penalty != 1.0:
	prev_tokens = generated[0].unique()
	penalty_logits = logits[:, prev_tokens]
	# For positive logits, divide by penalty; for negative, multiply.
	penalty_logits = torch.where(
	penalty_logits > 0,
	penalty_logits / repetition_penalty,
	penalty_logits * repetition_penalty,
	)
	logits[:, prev_tokens] = penalty_logits

	# Temperature scaling.
	logits = logits / temperature

	# Top-k filtering.
	if top_k > 0:
	top_k_values, _ = torch.topk(logits, top_k, dim=-1)
	threshold = top_k_values[:, -1].unsqueeze(-1)
	logits = logits.masked_fill(logits < threshold, float("-inf"))

	# Top-p (nucleus) filtering.
	if 0.0 < top_p < 1.0:
	sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
	cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

	# Mask tokens whose cumulative probability exceeds top_p.
	sorted_mask = cumulative_probs - F.softmax(sorted_logits, dim=-1) >= top_p
	sorted_logits = sorted_logits.masked_fill(sorted_mask, float("-inf"))

	# Scatter back to original ordering.
	logits = sorted_logits.scatter(1, sorted_indices, sorted_logits)

	# Sample from the filtered distribution.
	probs = F.softmax(logits, dim=-1)
	next_token = torch.multinomial(probs, num_samples=1) # (1, 1)

	generated = torch.cat([generated, next_token], dim=-1)

	if eos_token_id is not None and next_token.item() == eos_token_id:
	break

	return generated


	# ------------------------------------------------------------------ #
	# CLI #
	# ------------------------------------------------------------------ #

	def parse_args() -> argparse.Namespace:
	parser = argparse.ArgumentParser(
	description="BoccaccioAI text generation (free generation / Q&A)",
	)

	parser.add_argument(
	"--model-dir",
	type=str,
	required=True,
	help="Path to the saved model checkpoint directory.",
	)
	parser.add_argument(
	"--tokenizer-path",
	type=str,
	default="tokenizer/boccaccio-32k.json",
	help="Path to the tokenizer JSON file.",
	)
	parser.add_argument(
	"--config-path",
	type=str,
	default="configs/model.yaml",
	help="Path to the model config YAML file.",
	)
	parser.add_argument(
	"--config-variant",
	type=str,
	default="model",
	choices=["model", "nano"],
	help="Config variant to load ('model' or 'nano').",
	)
	parser.add_argument(
	"--mode",
	type=str,
	default="generate",
	choices=["generate", "qa", "chat"],
	help="Inference mode: 'generate' for free generation, 'qa' for Q&A, 'chat' for chat.",
	)
	parser.add_argument(
	"--prompt",
	type=str,
	default=None,
	help="Text prompt for generation/chat mode.",
	)
	parser.add_argument(
	"--system-prompt",
	type=str,
	default=None,
	help="Optional system prompt for chat mode.",
	)
	parser.add_argument(
	"--context",
	type=str,
	default=None,
	help="Context text for Q&A mode (legacy).",
	)
	parser.add_argument(
	"--question",
	type=str,
	default=None,
	help="Question text for Q&A mode (legacy).",
	)
	parser.add_argument(
	"--max-new-tokens",
	type=int,
	default=256,
	help="Maximum number of new tokens to generate.",
	)
	parser.add_argument(
	"--temperature",
	type=float,
	default=0.7,
	help="Sampling temperature.",
	)
	parser.add_argument(
	"--top-k",
	type=int,
	default=50,
	help="Top-k filtering value.",
	)
	parser.add_argument(
	"--top-p",
	type=float,
	default=0.9,
	help="Top-p (nucleus) sampling threshold.",
	)
	parser.add_argument(
	"--repetition-penalty",
	type=float,
	default=1.2,
	help="Repetition penalty (>1.0 reduces repetition, 1.0 = off).",
	)
	parser.add_argument(
	"--device",
	type=str,
	default="auto",
	choices=["cuda", "cpu", "auto"],
	help="Device to run inference on.",
	)

	return parser.parse_args()


	# ------------------------------------------------------------------ #
	# Main #
	# ------------------------------------------------------------------ #

	def main() -> None:
	args = parse_args()

	# ---- device -------------------------------------------------- #
	if args.device == "auto":
	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	else:
	device = torch.device(args.device)

	# ---- config -------------------------------------------------- #
	config = BoccaccioConfig.from_yaml(args.config_path, variant=args.config_variant)

	# ---- model --------------------------------------------------- #
	model = BoccaccioForCausalLM(config)

	checkpoint_dir = Path(args.model_dir)
	checkpoint_path = checkpoint_dir / "model.pt"
	state_dict = torch.load(checkpoint_path, map_location=device, weights_only=True)

	# Remap torch.compile prefix: _orig_mod.X -> X
	# Checkpoints saved with torch.compile wrap the model as _orig_mod.<inner>.
	cleaned = {}
	for key, value in state_dict.items():
	new_key = key
	if new_key.startswith("_orig_mod."):
	new_key = new_key[len("_orig_mod."):]
	cleaned[new_key] = value
	model.load_state_dict(cleaned)

	model.to(device)
	model.eval()

	# ---- tokenizer ----------------------------------------------- #
	tokenizer = Tokenizer.from_file(args.tokenizer_path)

	# ---- build prompt -------------------------------------------- #
	if args.mode == "chat":
	# Chat mode: use the same template as training (InstructionDataset).
	# Format: ### System:\n{system}\n### Utente:\n{prompt}\n### Assistente:\n
	if args.prompt is None:
	raise ValueError("Chat mode requires --prompt.")
	prompt = ""
	if args.system_prompt:
	prompt += f"### System:\n{args.system_prompt}\n"
	prompt += f"### Utente:\n{args.prompt}\n"
	prompt += f"### Assistente:\n"

	elif args.mode == "qa":
	# Legacy Q&A mode: map to chat format with context as system prompt.
	if args.context is None or args.question is None:
	raise ValueError("Q&A mode requires both --context and --question.")
	prompt = f"### System:\n{args.context}\n"
	prompt += f"### Utente:\n{args.question}\n"
	prompt += f"### Assistente:\n"

	else:
	# Free generation: raw prompt, no template.
	if args.prompt is None:
	raise ValueError("Generation mode requires --prompt.")
	prompt = args.prompt

	# ---- tokenize ------------------------------------------------ #
	encoding = tokenizer.encode(prompt)
	input_ids = torch.tensor([encoding.ids], dtype=torch.long, device=device)

	# ---- generate ------------------------------------------------ #
	eos_token_id = tokenizer.token_to_id("</s>")

	output_ids = generate(
	model,
	input_ids,
	max_new_tokens=args.max_new_tokens,
	temperature=args.temperature,
	top_k=args.top_k,
	top_p=args.top_p,
	repetition_penalty=args.repetition_penalty,
	eos_token_id=eos_token_id,
	)

	# ---- decode & print ------------------------------------------ #
	all_ids = output_ids.squeeze(0).tolist()
	prompt_len = len(encoding.ids)
	generated_ids = all_ids[prompt_len:]

	if args.mode in ("chat", "qa"):
	# In chat/qa mode, print only the generated response (not the prompt).
	response_text = tokenizer.decode(generated_ids).strip()
	print(response_text)
	else:
	# In generate mode, print the full text (prompt + generation).
	output_text = tokenizer.decode(all_ids)
	print(output_text)


	if __name__ == "__main__":
	main()