Publish preserved Nirnay 12K checkpoint with inference and evaluation evidence

855327f verified 4 days ago

3.63 kB

import os

import torch
from transformers import AutoTokenizer

os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")  # the tokenizer is shared with DataLoader worker processes

model_name = "Qwen/Qwen3-0.6B"
tokenizer = AutoTokenizer.from_pretrained(model_name)

OPTION_OPEN = tokenizer.encode("<option>")
OPTION_CLOSE = tokenizer.encode("</option>")
ANSWER_PROMPT = tokenizer.encode("The answer is:")

# Short natural-language task hint
TASK_PROMPTS = {
    "choice": "Choose the best option.",
    "score": "Rate it on the given scale.",
    "noul": "Answer yes or no.",
}


def truncate(text, max_tokens):
    tokens = tokenizer.encode(text)
    if len(tokens) <= max_tokens:
        return text
    return tokenizer.decode(tokens[:max_tokens])


def build_question(state, instructions, task_type, max_state_tokens):
    # Task hint and question come before the state, so every state token is read knowing what is asked
    return f"{TASK_PROMPTS[task_type]}\nQuestion: {instructions}\n\n{truncate(state, max_state_tokens)}\n"


def build_layout(question_ids, choice_ids, use_answer_prompt=True):
    """One sequence: question, then every choice wrapped in <option>...</option>, then "The answer is:".

    Returns token ids, position ids, a [L, L] attention mask (1 = may attend) and the read indices.
    Every choice starts at the same position and cannot see the other choices, so the choices are
    read in parallel: reordering them only reorders the outputs. The answer tokens come after the
    longest choice and see everything, which is where the choices are compared.
    """
    q = len(question_ids)
    input_ids = list(question_ids)
    position_ids = list(range(q))
    spans = []
    for ids in choice_ids:
        wrapped = OPTION_OPEN + ids + OPTION_CLOSE
        start = len(input_ids)
        input_ids += wrapped
        position_ids += range(q, q + len(wrapped))
        spans.append((start, len(input_ids)))  # [start, end)

    answer_start = len(input_ids)
    answer_pos = max(position_ids) + 1
    answer_ids = ANSWER_PROMPT if use_answer_prompt else []
    input_ids += answer_ids
    position_ids += range(answer_pos, answer_pos + len(answer_ids))
    L = len(input_ids)

    mask = torch.zeros((L, L))
    mask[:q, :q] = torch.tril(torch.ones(q, q))           # question: causal
    for start, end in spans:
        n = end - start
        mask[start:end, :q] = 1                            # each choice reads the question and state
        mask[start:end, start:end] = torch.tril(torch.ones(n, n))  # and itself, causally
    n = L - answer_start
    mask[answer_start:, :answer_start] = 1                 # answer reads everything before it
    mask[answer_start:, answer_start:] = torch.tril(torch.ones(n, n))

    # Choices are read at their last content token, not at the closing tag
    choice_read_idx = [end - len(OPTION_CLOSE) - 1 for _, end in spans]
    return {
        "input_ids": input_ids,
        "position_ids": position_ids,
        "attention_mask": mask,
        "choice_read_idx": choice_read_idx,
        "answer_end_idx": L - 1,
    }


def encode_example(state, instructions, choices, task_type, max_state_tokens, max_choice_tokens,
                   use_answer_prompt=True):
    question = build_question(state, instructions, task_type, max_state_tokens)
    choices = [truncate(choice, max_choice_tokens) for choice in choices]
    return build_layout(tokenizer.encode(question), [tokenizer.encode(choice) for choice in choices],
                        use_answer_prompt)

	import os

	import torch
	from transformers import AutoTokenizer

	os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") # the tokenizer is shared with DataLoader worker processes

	model_name = "Qwen/Qwen3-0.6B"
	tokenizer = AutoTokenizer.from_pretrained(model_name)

	OPTION_OPEN = tokenizer.encode("<option>")
	OPTION_CLOSE = tokenizer.encode("</option>")
	ANSWER_PROMPT = tokenizer.encode("The answer is:")

	# Short natural-language task hint
	TASK_PROMPTS = {
	"choice": "Choose the best option.",
	"score": "Rate it on the given scale.",
	"noul": "Answer yes or no.",
	}


	def truncate(text, max_tokens):
	tokens = tokenizer.encode(text)
	if len(tokens) <= max_tokens:
	return text
	return tokenizer.decode(tokens[:max_tokens])


	def build_question(state, instructions, task_type, max_state_tokens):
	# Task hint and question come before the state, so every state token is read knowing what is asked
	return f"{TASK_PROMPTS[task_type]}\nQuestion: {instructions}\n\n{truncate(state, max_state_tokens)}\n"


	def build_layout(question_ids, choice_ids, use_answer_prompt=True):
	"""One sequence: question, then every choice wrapped in <option>...</option>, then "The answer is:".

	Returns token ids, position ids, a [L, L] attention mask (1 = may attend) and the read indices.
	Every choice starts at the same position and cannot see the other choices, so the choices are
	read in parallel: reordering them only reorders the outputs. The answer tokens come after the
	longest choice and see everything, which is where the choices are compared.
	"""
	q = len(question_ids)
	input_ids = list(question_ids)
	position_ids = list(range(q))
	spans = []
	for ids in choice_ids:
	wrapped = OPTION_OPEN + ids + OPTION_CLOSE
	start = len(input_ids)
	input_ids += wrapped
	position_ids += range(q, q + len(wrapped))
	spans.append((start, len(input_ids))) # [start, end)

	answer_start = len(input_ids)
	answer_pos = max(position_ids) + 1
	answer_ids = ANSWER_PROMPT if use_answer_prompt else []
	input_ids += answer_ids
	position_ids += range(answer_pos, answer_pos + len(answer_ids))
	L = len(input_ids)

	mask = torch.zeros((L, L))
	mask[:q, :q] = torch.tril(torch.ones(q, q)) # question: causal
	for start, end in spans:
	n = end - start
	mask[start:end, :q] = 1 # each choice reads the question and state
	mask[start:end, start:end] = torch.tril(torch.ones(n, n)) # and itself, causally
	n = L - answer_start
	mask[answer_start:, :answer_start] = 1 # answer reads everything before it
	mask[answer_start:, answer_start:] = torch.tril(torch.ones(n, n))

	# Choices are read at their last content token, not at the closing tag
	choice_read_idx = [end - len(OPTION_CLOSE) - 1 for _, end in spans]
	return {
	"input_ids": input_ids,
	"position_ids": position_ids,
	"attention_mask": mask,
	"choice_read_idx": choice_read_idx,
	"answer_end_idx": L - 1,
	}


	def encode_example(state, instructions, choices, task_type, max_state_tokens, max_choice_tokens,
	use_answer_prompt=True):
	question = build_question(state, instructions, task_type, max_state_tokens)
	choices = [truncate(choice, max_choice_tokens) for choice in choices]
	return build_layout(tokenizer.encode(question), [tokenizer.encode(choice) for choice in choices],
	use_answer_prompt)