simonbb 0.0.7__tar.gz → 0.0.8__tar.gz

{simonbb-0.0.7 → simonbb-0.0.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: simonbb
-Version: 0.0.7
+Version: 0.0.8
 Summary: A small package to facilitate teaching AI courses at business schools
 Project-URL: Homepage, https://simonbb.com
 Author-email: Huaxia Rui <huaxia.rui@simon.rochester.edu>
@@ -10,12 +10,6 @@ Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.9
 Requires-Dist: gymnasium~=1.3.0
-Requires-Dist: langchain-chroma~=1.1.0
-Requires-Dist: langchain-community~=0.4.1
-Requires-Dist: langchain-core~=1.2.23
-Requires-Dist: langchain-huggingface~=1.2.1
-Requires-Dist: langchain~=1.2.13
-Requires-Dist: langgraph~=1.1.3
 Requires-Dist: matplotlib~=3.10.8
 Requires-Dist: numpy~=2.4.0
 Requires-Dist: opencv-python~=5.0.0.93

{simonbb-0.0.7 → simonbb-0.0.8}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "simonbb"
-version = "0.0.7"
+version = "0.0.8"
 authors = [
 	{ name="Huaxia Rui", email="huaxia.rui@simon.rochester.edu" },
 ]
@@ -23,12 +23,6 @@ dependencies = [
 	"torchvision~=0.24.1",
 	"numpy~=2.4.0",
 	"transformers~=4.57.3",
-	"langchain-core~=1.2.23",
-	"langchain~=1.2.13",
-	"langgraph~=1.1.3",
-	"langchain-community~=0.4.1",
-	"langchain-chroma~=1.1.0",
-	"langchain-huggingface~=1.2.1",
 	"sentence-transformers~=5.3.0",
 	"gymnasium~=1.3.0",
 	"pytorch-ignite~=0.5.4",

{simonbb-0.0.7 → simonbb-0.0.8}/src/simonbb/cis433/experience.py

@@ -31,9 +31,7 @@ class ExperienceSource:
     """
     Item = tt.Tuple[Experience, ...]
 
-    def __init__(self, env: gym.Env | tt.Collection[gym.Env], agent: BaseAgent,
-                 steps_count: int = 2, steps_delta: int = 1,
-                 env_seed: tt.Optional[int] = None):
+    def __init__(self, env: gym.Env | tt.Collection[gym.Env], agent: BaseAgent, steps_count: int = 2, steps_delta: int = 1, env_seed: tt.Optional[int] = None):
         """
         Create simple experience source
         :param env: environment or list of environments to be used
@@ -46,7 +44,7 @@ class ExperienceSource:
         if isinstance(env, (list, tuple)):
             self.pool = env
             # do the check for the multiple copies passed
-            ids = set(id(e) for e in env)
+            ids = set(id(e) for e in env) # id() returns a unique, constant integer identifier for a specified object during its lifetime
             if len(ids) < len(env):
                 raise ValueError("You passed single environment instance multiple times")
         else:

simonbb-0.0.8/src/simonbb/cis433/math500.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import re
+import time
+import json
+from pathlib import Path
+from ..utils import grade_math_answer, eta_progress_message
+
+RE_NUMBER = re.compile( r"-?(?:\d+/\d+|\d+(?:\.\d+)?(?:[eE][+-]?\d+)?)" )
+def get_last_boxed(text):
+    # Find the last occurrence of "\boxed"
+    boxed_start_idx = text.rfind(r"\boxed")
+    if boxed_start_idx == -1:
+        return None
+    # Get position after "\boxed"
+    current_idx = boxed_start_idx + len(r"\boxed")
+    # Skip any whitespace after "\boxed"
+    while current_idx < len(text) and text[current_idx].isspace():
+        current_idx += 1
+    # Expect an opening brace "{"
+    if current_idx >= len(text) or text[current_idx] != "{":
+        return None
+    # Parse the braces with nesting
+    current_idx += 1
+    brace_depth = 1
+    content_start_idx = current_idx
+    while current_idx < len(text) and brace_depth > 0:
+        char = text[current_idx]
+        if char == "{":
+            brace_depth += 1
+        elif char == "}":
+            brace_depth -= 1
+        current_idx += 1
+    # Account for unbalanced braces
+    if brace_depth != 0:
+        return None
+    # Extract content inside the outermost braces
+    return text[content_start_idx:current_idx-1]
+
+def extract_final_candidate(text, fallback="number_then_full"):
+    # Default return value if nothing matches
+    result = ""
+    if text:        
+        boxed = get_last_boxed(text.strip()) # Prefer the last boxed expression if present
+        if boxed:
+            result = boxed.strip().strip("$ ")
+        # If no boxed expression, try fallback
+        elif fallback in ("number_then_full", "number_only"):
+            m = RE_NUMBER.findall(text)
+            if m: # Use last number
+                result = m[-1]
+            elif fallback == "number_then_full":
+                # Else return full text if no number found
+                result = text
+    return result
+
+def render_prompt(prompt):
+    template = (
+        "You are a helpful math assistant.\n"
+        "Answer the question and write the final result on a new line as:\n"
+        "\\boxed{ANSWER}\n\n"
+        f"Question:\n{prompt}\n\nAnswer:"
+    )
+    return template
+
+
+def evaluate_math500_stream( model, tokenizer, device, generate_func, math_data, out_path=None, max_new_tokens=512, verbose=False ):
+    num_examples = len(math_data)
+    num_correct = 0
+    start_time = time.time()
+    with open(Path( out_path ), "w", encoding="utf-8") as f:  # Save results for inspection
+        for i, row in enumerate(math_data, start=1):
+            prompt = render_prompt(row["problem"])    # 1. Apply prompt template
+            gen_text = generate_func( model=model, tokenizer=tokenizer, prompt=prompt, device=device, max_new_tokens=max_new_tokens, verbose=verbose )
+            extracted = extract_final_candidate( gen_text )
+            is_correct = grade_math_answer( extracted, row["answer"] )
+            num_correct += int(is_correct)
+            record = {  # Record to be saved for inspection
+                "index": i,
+                "problem": row["problem"],
+                "gtruth_answer": row["answer"],
+                "generated_text": gen_text,
+                "extracted": extracted,
+                "correct": bool(is_correct),
+            }
+            f.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+            progress_msg = eta_progress_message( processed=i, total=num_examples, start_time=start_time, show_eta=True, label="MATH-500" )
+            print(progress_msg, end="\r", flush=True)
+            if verbose:  # Print responses during the generation process
+                print(
+                    f"\n\n{'='*50}\n{progress_msg}\n"
+                    f"{'='*50}\nExtracted: {extracted}\n"
+                    f"Expected:  {row['answer']}\n"
+                    f"Correct so far: {num_correct}\n{'-'*50}"
+                )
+
+    # Print summary information
+    seconds_elapsed = time.time() - start_time
+    acc = num_correct / num_examples if num_examples else 0.0
+    print(f"\nAccuracy: {acc*100:.1f}% ({num_correct}/{num_examples})")
+    print(f"Total time: {seconds_elapsed/60:.1f} min")
+    print(f"Logs written to: {out_path}")
+    return num_correct, num_examples, acc

simonbb-0.0.8/src/simonbb/cis433/qwen3.py

@@ -0,0 +1,639 @@
+"""
+This code is modified based on Sebastian Raschka's Implementation for personal use. No copyright claimed.
+Original Code Repository: https://github.com/rasbt/reasoning-from-scratch
+Original model download at 
+https://huggingface.co/rasbt/qwen3-from-scratch/resolve/main/qwen3-0.6B-base.pth
+https://huggingface.co/rasbt/qwen3-from-scratch/resolve/main/qwen3-0.6B-reasoning.pth
+https://huggingface.co/rasbt/qwen3-from-scratch/resolve/main/tokenizer-base.json
+https://huggingface.co/rasbt/qwen3-from-scratch/resolve/main/tokenizer-reasoning.json
+"""
+
+import json
+import os
+import re
+from pathlib import Path
+from safetensors.torch import load_file 
+
+import torch
+import torch.nn as nn
+
+# 0.6 billion parameters
+QWEN_CONFIG_06_B = {
+    "vocab_size": 151_936,     # Vocabulary size
+    "context_length": 40_960,  # Length originally used during training
+    "emb_dim": 1024,           # Embedding dimension
+    "n_heads": 16,             # Number of attention heads
+    "n_layers": 28,            # Number of layers
+    "hidden_dim": 3072,        # Size of intermediate dim in FeedForward
+    "head_dim": 128,           # Size of the heads in GQA
+    "qk_norm": True,           # Whether to normalize queries & keys in GQA
+    "n_kv_groups": 8,          # Key-Value groups for GQA
+    "rope_base": 1_000_000.0,  # The base in RoPE's "theta"
+    "dtype": torch.bfloat16,   # Lower-precision dtype to reduce memory
+}
+
+QWEN3_CONFIG_1_7B = {
+    "vocab_size": 151_936,
+    "context_length": 40_960,
+    "emb_dim": 2048,                 # 2x larger than 0.6B model
+    "n_heads": 16,
+    "n_layers": 28,
+    "hidden_dim": 6144,              # 2x larger than 0.6B model
+    "head_dim": 128,
+    "qk_norm": True,
+    "n_kv_groups": 8,
+    "rope_base": 1_000_000.0,
+    "dtype": torch.bfloat16,
+}
+
+# 4 billion parameters
+QWEN3_CONFIG_4B = {
+    "vocab_size": 151_936,
+    "context_length": 40_960,
+    "emb_dim": 2560,                 # 25% larger than above
+    "n_heads": 32,                   # 2x larger than above
+    "n_layers": 36,                  # 29% larger than above
+    "hidden_dim": 9728,              # ~3x larger than above
+    "head_dim": 128,
+    "qk_norm": True,
+    "n_kv_groups": 8,
+    "rope_base": 1_000_000.0,
+    "dtype": torch.bfloat16,
+}
+
+# 8 billion parameters
+QWEN3_CONFIG_8B = {
+    "vocab_size": 151_936,
+    "context_length": 40_960,
+    "emb_dim": 4096,                 # 60% larger than above
+    "n_heads": 32,
+    "n_layers": 36,                  # 26% larger than above
+    "hidden_dim": 12288,
+    "head_dim": 128,
+    "qk_norm": True,
+    "n_kv_groups": 8,
+    "rope_base": 1_000_000.0,
+    "dtype": torch.bfloat16,
+}
+
+# 14 billion parameters
+QWEN3_CONFIG_14B = {
+        "vocab_size": 151_936,
+        "context_length": 40_960,
+        "emb_dim": 5120,                 # 25% larger than above
+        "n_heads": 40,                   # 25% larger than above
+        "n_layers": 40,                  # 11% larger than above
+        "hidden_dim": 17408,             # 42% larger than above
+        "head_dim": 128,
+        "qk_norm": True,
+        "n_kv_groups": 8,
+        "rope_base": 1_000_000.0,
+        "dtype": torch.bfloat16,
+}
+
+QWEN3_CONFIG_32B = {
+        "vocab_size": 151_936,
+        "context_length": 40_960,
+        "emb_dim": 5120,
+        "n_heads": 64,                   # 60% larger than above
+        "n_layers": 64,                  # 60% larger than above
+        "hidden_dim": 25600,             # 47% larger than above
+        "head_dim": 128,
+        "qk_norm": True,
+        "n_kv_groups": 8,
+        "rope_base": 1_000_000.0,
+        "dtype": torch.bfloat16,
+}
+
+class Qwen3Model(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+
+        # Main model parameters
+        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"], dtype=cfg["dtype"])
+
+        self.trf_blocks = nn.ModuleList(  # ModuleList since Sequential can only accept one input, and we need `x, mask, cos, sin`
+            [TransformerBlock(cfg) for _ in range(cfg["n_layers"])]
+        )
+        self.final_norm = RMSNorm(cfg["emb_dim"])
+        self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False, dtype=cfg["dtype"])
+
+        # Reusable utilities
+        if cfg["head_dim"] is None:
+            head_dim = cfg["emb_dim"] // cfg["n_heads"]
+        else:
+            head_dim = cfg["head_dim"]
+        cos, sin = compute_rope_params(
+            head_dim=head_dim,
+            theta_base=cfg["rope_base"],
+            context_length=cfg["context_length"]
+        )
+        self.register_buffer("cos", cos, persistent=False)
+        self.register_buffer("sin", sin, persistent=False)
+        self.cfg = cfg
+        self.current_pos = 0  # Track current position in KV cache
+
+    def forward(self, in_idx, cache=None):
+        # Forward pass
+        tok_embeds = self.tok_emb(in_idx)
+        x = tok_embeds
+
+        num_tokens = x.shape[1]
+        if cache is not None:
+            pos_start = self.current_pos
+            pos_end = pos_start + num_tokens
+            self.current_pos = pos_end
+            mask = torch.triu(
+                torch.ones(pos_end, pos_end, device=x.device, dtype=torch.bool), diagonal=1
+            )[pos_start:pos_end, :pos_end]
+        else:
+            pos_start = 0  # Not strictly necessary but helps torch.compile
+            mask = torch.triu(
+                torch.ones(num_tokens, num_tokens, device=x.device, dtype=torch.bool), diagonal=1
+            )
+        # Prefill (no cache): mask starts as (num_tokens, num_tokens)
+        # Cached decoding: mask starts as (num_tokens, prev_k_number_tokens + num_tokens)
+        #
+        # We add two leading dimensions so the mask becomes
+        # (1, 1, num_tokens, num_tokens) during prefill and
+        # (1, 1, num_tokens, total_key_tokens) during cached decoding.
+        # These extra dimensions let PyTorch broadcast the same mask
+        # across all batches and attention heads when applying it to
+        # attn_scores of shape (batch, num_heads, num_tokens, total_key_tokens).
+        mask = mask[None, None, :, :]  # broadcast mask
+
+        for i, block in enumerate(self.trf_blocks):
+            blk_cache = cache.get(i) if cache else None
+            x, new_blk_cache = block(x, mask, self.cos, self.sin,
+                                     start_pos=pos_start,
+                                     cache=blk_cache)
+            if cache is not None:
+                cache.update(i, new_blk_cache)
+
+        x = self.final_norm(x)
+        logits = self.out_head(x.to(self.cfg["dtype"]))
+        return logits
+
+    def reset_kv_cache(self):
+        self.current_pos = 0
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.att = GroupedQueryAttention(
+            d_in=cfg["emb_dim"],
+            num_heads=cfg["n_heads"],
+            head_dim=cfg["head_dim"],
+            num_kv_groups=cfg["n_kv_groups"],
+            qk_norm=cfg["qk_norm"],
+            dtype=cfg["dtype"]
+        )
+        self.ff = FeedForward(cfg)
+        self.norm1 = RMSNorm(cfg["emb_dim"], eps=1e-6)
+        self.norm2 = RMSNorm(cfg["emb_dim"], eps=1e-6)
+
+    def forward(self, x, mask, cos, sin, start_pos=0, cache=None):
+        # Shortcut connection for attention block
+        shortcut = x
+        x = self.norm1(x)
+        x, next_cache = self.att(x, mask, cos, sin, start_pos=start_pos, cache=cache)  # Shape [batch_size, num_tokens, emb_size]
+        x = x + shortcut  # Add the original input back
+
+        # Shortcut connection for feed-forward block
+        shortcut = x
+        x = self.norm2(x)
+        x = self.ff(x)
+        x = x + shortcut  # Add the original input back
+
+        return x, next_cache
+
+
+class FeedForward(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.fc1 = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], dtype=cfg["dtype"], bias=False)
+        self.fc2 = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], dtype=cfg["dtype"], bias=False)
+        self.fc3 = nn.Linear(cfg["hidden_dim"], cfg["emb_dim"], dtype=cfg["dtype"], bias=False)
+
+    def forward(self, x):
+        x_fc1 = self.fc1(x)
+        x_fc2 = self.fc2(x)
+        x = nn.functional.silu(x_fc1) * x_fc2
+        return self.fc3(x)
+
+
+class GroupedQueryAttention(nn.Module):
+    def __init__(
+        self, d_in, num_heads, num_kv_groups, head_dim=None, qk_norm=False, dtype=None
+    ):
+        super().__init__()
+        assert num_heads % num_kv_groups == 0, "num_heads must be divisible by num_kv_groups"
+
+        self.num_heads = num_heads
+        self.num_kv_groups = num_kv_groups
+        self.group_size = num_heads // num_kv_groups
+
+        if head_dim is None:
+            assert d_in % num_heads == 0, "`d_in` must be divisible by `num_heads` if `head_dim` is not set"
+            head_dim = d_in // num_heads
+
+        self.head_dim = head_dim
+        self.d_out = num_heads * head_dim
+
+        self.W_query = nn.Linear(d_in, self.d_out, bias=False, dtype=dtype)
+        self.W_key = nn.Linear(d_in, num_kv_groups * head_dim, bias=False, dtype=dtype)
+        self.W_value = nn.Linear(d_in, num_kv_groups * head_dim, bias=False, dtype=dtype)
+
+        self.out_proj = nn.Linear(self.d_out, d_in, bias=False, dtype=dtype)
+
+        if qk_norm:
+            self.q_norm = RMSNorm(head_dim, eps=1e-6)
+            self.k_norm = RMSNorm(head_dim, eps=1e-6)
+        else:
+            self.q_norm = self.k_norm = None
+
+    def forward(self, x, mask, cos, sin, start_pos=0, cache=None):
+        b, num_tokens, _ = x.shape
+
+        # Apply projections
+        queries = self.W_query(x)  # (b, num_tokens, num_heads * head_dim)
+        keys = self.W_key(x)       # (b, num_tokens, num_kv_groups * head_dim)
+        values = self.W_value(x)   # (b, num_tokens, num_kv_groups * head_dim)
+
+        # Reshape to heads / kv-groups
+        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim).transpose(1, 2)
+        keys_new = keys.view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)
+        values_new = values.view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)
+
+        # Optional normalization
+        if self.q_norm:
+            queries = self.q_norm(queries)
+        if self.k_norm:
+            keys_new = self.k_norm(keys_new)
+
+        # Apply RoPE
+        queries = apply_rope(queries, cos, sin, offset=start_pos)
+        keys_new = apply_rope(keys_new, cos, sin, offset=start_pos)
+
+        if cache is not None:
+            prev_k, prev_v = cache
+            keys = torch.cat([prev_k, keys_new], dim=2)
+            values = torch.cat([prev_v, values_new], dim=2)
+        else:
+            start_pos = 0  # reset RoPE
+            keys, values = keys_new, values_new
+        next_cache = (keys, values)
+
+        # Expand K and V to match number of heads
+        keys = keys.repeat_interleave(self.group_size, dim=1)
+        values = values.repeat_interleave(self.group_size, dim=1)
+
+        # Attention
+        attn_scores = queries @ keys.transpose(2, 3)
+        attn_scores = attn_scores.masked_fill(mask, -torch.inf)
+        attn_weights = torch.softmax(attn_scores / self.head_dim**0.5, dim=-1)
+
+        context = (attn_weights @ values).transpose(1, 2).reshape(b, num_tokens, self.d_out)
+        return self.out_proj(context), next_cache
+
+
+# ==============================================================================
+# RoPE implementation summary
+#
+#
+# There are two common styles to implement RoPE, which are
+# mathematically equivalent;
+# they mainly differ in how the rotation matrix pairs dimensions.
+#
+# 1) Split-halves style (this repo, Hugging Face Transformers):
+#
+#   For hidden dim d = 4 (example):
+#
+#       [ x0   x1 | x2   x3 ]
+#         │    │    │    │
+#         ▼    ▼    ▼    ▼
+#        cos  cos  sin  sin
+#
+#   Rotation matrix:
+#
+#       [ cosθ0   0    -sinθ0   0   ]
+#       [  0    cosθ1    0    -sinθ1]
+#       [ sinθ0   0     cosθ0   0   ]
+#       [  0    sinθ1    0     cosθ1]
+#
+#   Here, the embedding dims are split into two halves and then
+#   each one is rotated in blocks.
+#
+#
+# 2) Interleaved (even/odd) style (original paper, Llama repo):
+#
+#   For hidden dim d = 4 (example):
+#
+#       [ x0   x1   x2   x3 ]
+#         │    │    │    │
+#         ▼    ▼    ▼    ▼
+#        cos  sin  cos  sin
+#
+#   Rotation matrix:
+#
+#       [ cosθ0  -sinθ0   0       0    ]
+#       [ sinθ0   cosθ0   0       0    ]
+#       [  0        0    cosθ1  -sinθ1 ]
+#       [  0        0    sinθ1   cosθ1 ]
+#
+#
+#   Here, embedding dims are interleaved as even/odd cosine/sine pairs.
+#
+# Both layouts encode the same relative positions; the only difference is how
+# dimensions are paired.
+# ==============================================================================
+
+
+def compute_rope_params(head_dim, theta_base=10_000, context_length=4096, dtype=torch.float32):
+    assert head_dim % 2 == 0, "Embedding dimension must be even"
+
+    # Compute the inverse frequencies
+    inv_freq = 1.0 / (theta_base ** (torch.arange(0, head_dim, 2, dtype=dtype)[: (head_dim // 2)].float() / head_dim))
+
+    # Generate position indices
+    positions = torch.arange(context_length, dtype=dtype)
+
+    # Compute the angles
+    angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0)  # Shape: (context_length, head_dim // 2)
+
+    # Expand angles to match the head_dim
+    angles = torch.cat([angles, angles], dim=1)  # Shape: (context_length, head_dim)
+
+    # Precompute sine and cosine
+    cos = torch.cos(angles)
+    sin = torch.sin(angles)
+
+    return cos, sin
+
+
+def apply_rope(x, cos, sin, offset=0):
+    # x: (batch_size, num_heads, seq_len, head_dim)
+    batch_size, num_heads, seq_len, head_dim = x.shape
+    assert head_dim % 2 == 0, "Head dimension must be even"
+
+    # Split x into first half and second half
+    x1 = x[..., : head_dim // 2]  # First half
+    x2 = x[..., head_dim // 2:]  # Second half
+
+    # Adjust sin and cos shapes
+    cos = cos[offset:offset + seq_len, :].unsqueeze(0).unsqueeze(0)  # Shape: (1, 1, seq_len, head_dim)
+    sin = sin[offset:offset + seq_len, :].unsqueeze(0).unsqueeze(0)
+
+    # Apply the rotary transformation
+    rotated = torch.cat((-x2, x1), dim=-1)
+    x_rotated = (x * cos) + (rotated * sin)
+
+    # It's ok to use lower-precision after applying cos and sin rotation
+    return x_rotated.to(dtype=x.dtype)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, emb_dim, eps=1e-6, bias=False, qwen3_compatible=True):
+        super().__init__()
+        self.eps = eps
+        self.qwen3_compatible = qwen3_compatible
+        self.scale = nn.Parameter(torch.ones(emb_dim))
+        self.shift = nn.Parameter(torch.zeros(emb_dim)) if bias else None
+
+    def forward(self, x):
+        input_dtype = x.dtype
+
+        if self.qwen3_compatible:
+            x = x.to(torch.float32)
+
+        variance = x.pow(2).mean(dim=-1, keepdim=True)
+        norm_x = x * torch.rsqrt(variance + self.eps)
+        norm_x = norm_x * self.scale
+
+        if self.shift is not None:
+            norm_x = norm_x + self.shift
+
+        return norm_x.to(input_dtype)
+
+
+class Qwen3Tokenizer:
+    _SPECIALS = [
+        "<|endoftext|>",
+        "<|im_start|>", "<|im_end|>",
+        "<|object_ref_start|>", "<|object_ref_end|>",
+        "<|box_start|>", "<|box_end|>",
+        "<|quad_start|>", "<|quad_end|>",
+        "<|vision_start|>", "<|vision_end|>",
+        "<|vision_pad|>", "<|image_pad|>", "<|video_pad|>",
+    ]
+    _SPLIT_RE = re.compile(r"(<\|[^>]+?\|>)")
+
+    def __init__(self, tokenizer_file_path="tokenizer-base.json",
+                 apply_chat_template=False,
+                 add_generation_prompt=False,
+                 add_thinking=False):
+        from tokenizers import Tokenizer
+
+        self.apply_chat_template = apply_chat_template
+        self.add_generation_prompt = add_generation_prompt
+        self.add_thinking = add_thinking
+
+        tok_path = Path(tokenizer_file_path)
+        if not tok_path.is_file():
+            raise FileNotFoundError(
+                f"Tokenizer file '{tok_path}' not found. Please ensure it's available."
+            )
+
+        self._tok = Tokenizer.from_file(str(tok_path))
+        self._special_to_id = {t: self._tok.token_to_id(t) for t in self._SPECIALS}
+
+        self.pad_token = "<|endoftext|>"
+        self.pad_token_id = self._special_to_id.get(self.pad_token)
+
+        # Match HF behavior: chat model → <|im_end|>, base model → <|endoftext|>
+        fname = tok_path.name.lower()
+        if "base" in fname and "reasoning" not in fname:
+            self.eos_token = "<|endoftext|>"
+        else:
+            self.eos_token = "<|im_end|>"
+        self.eos_token_id = self._special_to_id.get(self.eos_token)
+
+    def encode(self, prompt, chat_wrapped=None):
+        if chat_wrapped is None:
+            chat_wrapped = self.apply_chat_template
+
+        stripped = prompt.strip()
+        if stripped in self._special_to_id and "\n" not in stripped:
+            return [self._special_to_id[stripped]]
+
+        if chat_wrapped:
+            prompt = self._wrap_chat(prompt)
+
+        ids = []
+        for part in filter(None, self._SPLIT_RE.split(prompt)):
+            if part in self._special_to_id:
+                ids.append(self._special_to_id[part])
+            else:
+                ids.extend(self._tok.encode(part).ids)
+        return ids
+
+    def decode(self, token_ids):
+        return self._tok.decode(token_ids, skip_special_tokens=False)
+
+    def _wrap_chat(self, user_msg):
+        s = f"<|im_start|>user\n{user_msg}<|im_end|>\n"
+        if self.add_generation_prompt:
+            s += "<|im_start|>assistant"
+            if self.add_thinking:
+                s += "\n"  # insert no <think> tag, just a new line
+            else:
+                s += "\n<think>\n\n</think>\n\n"
+        return s
+
+def load_hf_weights( local_dir ):
+    index_path = os.path.join(local_dir, "model.safetensors.index.json")
+    single_file_path = os.path.join(local_dir, "model.safetensors")
+    if os.path.exists(index_path):  # Multi-shard model
+        with open(index_path, "r") as f:
+            index = json.load(f)
+        weights_dict = {}
+        for filename in set(index["weight_map"].values()):
+            shard_path = os.path.join(local_dir, filename)
+            shard = load_file(shard_path)
+            weights_dict.update(shard)
+    elif os.path.exists(single_file_path): # Single-shard model        
+        weights_dict = load_file(single_file_path)
+    else:
+        raise FileNotFoundError("No model.safetensors or model.safetensors.index.json found.")
+    return weights_dict
+
+def load_hf_weights_into_qwen(model, param_config, params):
+    """
+    Only used in Appendix D for loading the other Qwen3 variants.
+    """
+    def assign(left, right, tensor_name="unknown"):
+        if left.shape != right.shape:
+            raise ValueError(f"Shape mismatch in tensor '{tensor_name}'. Left: {left.shape}, Right: {right.shape}")
+
+        with torch.no_grad():
+            if isinstance(right, torch.Tensor):
+                left.copy_(right)
+            else:
+                left.copy_(torch.as_tensor(right, dtype=left.dtype, device=left.device))
+
+        return left
+
+    model.tok_emb.weight = assign(model.tok_emb.weight, params["model.embed_tokens.weight"], "model.embed_tokens.weight")
+
+    for l in range(param_config["n_layers"]):  # noqa: E741
+        block = model.trf_blocks[l]
+        att = block.att
+
+        # Q, K, V projections
+        att.W_query.weight = assign(
+            att.W_query.weight,
+            params[f"model.layers.{l}.self_attn.q_proj.weight"],
+            f"model.layers.{l}.self_attn.q_proj.weight"
+        )
+        att.W_key.weight = assign(
+            att.W_key.weight,
+            params[f"model.layers.{l}.self_attn.k_proj.weight"],
+            f"model.layers.{l}.self_attn.k_proj.weight"
+        )
+        att.W_value.weight = assign(
+            att.W_value.weight,
+            params[f"model.layers.{l}.self_attn.v_proj.weight"],
+            f"model.layers.{l}.self_attn.v_proj.weight"
+        )
+
+        # Output projection
+        att.out_proj.weight = assign(
+            att.out_proj.weight,
+            params[f"model.layers.{l}.self_attn.o_proj.weight"],
+            f"model.layers.{l}.self_attn.o_proj.weight"
+        )
+
+        # QK norms
+        if hasattr(att, "q_norm") and att.q_norm is not None:
+            att.q_norm.scale = assign(
+                att.q_norm.scale,
+                params[f"model.layers.{l}.self_attn.q_norm.weight"],
+                f"model.layers.{l}.self_attn.q_norm.weight"
+            )
+        if hasattr(att, "k_norm") and att.k_norm is not None:
+            att.k_norm.scale = assign(
+                att.k_norm.scale,
+                params[f"model.layers.{l}.self_attn.k_norm.weight"],
+                f"model.layers.{l}.self_attn.k_norm.weight"
+            )
+
+        # Attention layernorm
+        block.norm1.scale = assign(
+            block.norm1.scale,
+            params[f"model.layers.{l}.input_layernorm.weight"],
+            f"model.layers.{l}.input_layernorm.weight"
+        )
+
+        # Feedforward weights
+        if "num_experts" in param_config:
+            # Load router (gating) weights
+            block.ff.gate.weight = assign(
+                block.ff.gate.weight,
+                params[f"model.layers.{l}.mlp.gate.weight"],
+                f"model.layers.{l}.mlp.gate.weight"
+            )
+            # Load expert weights
+            for e in range(param_config["num_experts"]):
+                prefix = f"model.layers.{l}.mlp.experts.{e}"
+                block.ff.fc1[e].weight = assign(
+                    block.ff.fc1[e].weight,
+                    params[f"{prefix}.gate_proj.weight"],
+                    f"{prefix}.gate_proj.weight"
+                )
+                block.ff.fc2[e].weight = assign(
+                    block.ff.fc2[e].weight,
+                    params[f"{prefix}.up_proj.weight"],
+                    f"{prefix}.up_proj.weight"
+                )
+                block.ff.fc3[e].weight = assign(
+                    block.ff.fc3[e].weight,
+                    params[f"{prefix}.down_proj.weight"],
+                    f"{prefix}.down_proj.weight"
+                )
+                # After assigning weights, move the expert layers from meta to CPU
+                block.ff.fc1[e] = block.ff.fc1[e].to("cpu")
+                block.ff.fc2[e] = block.ff.fc2[e].to("cpu")
+                block.ff.fc3[e] = block.ff.fc3[e].to("cpu")
+
+        else:
+            block.ff.fc1.weight = assign(
+                block.ff.fc1.weight,
+                params[f"model.layers.{l}.mlp.gate_proj.weight"],
+                f"model.layers.{l}.mlp.gate_proj.weight"
+            )
+            block.ff.fc2.weight = assign(
+                block.ff.fc2.weight,
+                params[f"model.layers.{l}.mlp.up_proj.weight"],
+                f"model.layers.{l}.mlp.up_proj.weight"
+            )
+            block.ff.fc3.weight = assign(
+                block.ff.fc3.weight,
+                params[f"model.layers.{l}.mlp.down_proj.weight"],
+                f"model.layers.{l}.mlp.down_proj.weight"
+            )
+
+        block.norm2.scale = assign(
+            block.norm2.scale,
+            params[f"model.layers.{l}.post_attention_layernorm.weight"],
+            f"model.layers.{l}.post_attention_layernorm.weight"
+        )
+
+    # Final normalization and output head
+    model.final_norm.scale = assign(model.final_norm.scale, params["model.norm.weight"], "model.norm.weight")
+
+    if "lm_head.weight" in params:
+        model.out_head.weight = assign(model.out_head.weight, params["lm_head.weight"], "lm_head.weight")
+    else:
+        model.out_head.weight = model.tok_emb.weight
+        print("Model uses weight tying.")

{simonbb-0.0.7 → simonbb-0.0.8}/src/simonbb/cis433/utils.py

@@ -14,6 +14,37 @@ import torch.nn as nn
 
 #%%
 
+def top_p_filter(probas, top_p):
+    if top_p is None or top_p >= 1.0:
+        return probas    
+    sorted_probas, sorted_idx = torch.sort(probas, dim=1, descending=True)    
+    cumprobas = torch.cumsum(sorted_probas, dim=1)    
+    cumprobas_before_atom = cumprobas - sorted_probas
+    keep = cumprobas_before_atom < top_p    
+    keep[:, 0] = True # Always keep at least one token (fallback for very small/non-positive top_p)
+    kept_sorted = torch.where( keep, sorted_probas, torch.zeros_like(sorted_probas) )    
+    filtered = torch.zeros_like(probas).scatter(1, sorted_idx, kept_sorted) # Map back to original order    
+    denom = torch.sum(filtered, dim=1, keepdim=True).clamp_min(1e-12) # Renormalize to sum to 1
+    return filtered / denom
+
+class KVCache:
+    def __init__(self, n_layers):
+        self.cache = [None] * n_layers
+
+    def get(self, layer_idx):
+        return self.cache[layer_idx]
+
+    def update(self, layer_idx, value):
+        self.cache[layer_idx] = value
+
+    def get_all(self):
+        return self.cache
+
+    def reset(self):
+        for i in range(len(self.cache)):
+            self.cache[i] = None
+            
+            
 class ModelCheckpoint:
     def __init__(self, filepath, monitor="val_loss", save_best_only=True, save_optimizer_state=False):
         self.filepath = filepath
@@ -69,6 +100,25 @@ def stats_per_channel(loader):
     norm_std = total_stds / len(loader.dataset) # (1, n_channels)
     return norm_mean.squeeze(), norm_std.squeeze() # without squeezing, runtime error: not all boolean value of tensor with more than one value is ambiguous
 
+import warnings
+def generate_stats(output_token_ids, tokenizer, start_time, end_time):
+    total_time = end_time - start_time
+    print(f"\n\nTime: {total_time:.2f} sec")
+    print(f"{int(output_token_ids.numel() / total_time)} tokens/sec")
+    for name, backend in (("CUDA", getattr(torch, "cuda", None)), ("XPU", getattr(torch, "xpu", None))):
+        if backend is not None and backend.is_available():
+            # Check whether we are actually using this backend
+            device_type = output_token_ids.device.type
+            if device_type != name.lower():
+                warnings.warn( f"{name} is available but tensors are on {device_type}. Memory stats may be 0." )
+                # Synchronize if supported (important for async backends)
+            if hasattr(backend, "synchronize"):
+                backend.synchronize()            
+            max_mem_bytes = backend.max_memory_allocated()
+            max_mem_gb = max_mem_bytes / (1024 ** 3)
+            print(f"Max {name} memory allocated: {max_mem_gb:.2f} GB")
+            backend.reset_peak_memory_stats()
+            
 #%% Train Loop: tested on classification, regression, image segmentation, translation, gpt2
 
 def train( model, train_dataloader, val_dataloader, optimizer, scheduler=None, clip_grad=None, loss_fn=nn.CrossEntropyLoss(), callbacks=None, device=None, evaluation=True, n_epochs=1, n_batch_per_report=5, accuracy_fn=None, lr_history=None ):
@@ -311,6 +361,95 @@ def evaluate(model, loader, device, loss_fn=nn.CrossEntropyLoss(), accuracy_fn=N
         val_accuracy = val_accuracy.item() # extract the value of a single-item tensor 
     return {'loss':val_loss, 'accuracy':val_accuracy}
 
+#%%
+
+@torch.inference_mode()
+def generate_text_basic_stream( model, token_ids, max_new_tokens, eos_token_id=None, kv_cache = True ):
+    model.eval()
+    if kv_cache:
+        cache = KVCache(n_layers=model.cfg["n_layers"])
+        model.reset_kv_cache()
+        out = model(token_ids, cache=cache)[:, -1]
+    else:
+        out = model(token_ids)[:, -1]
+    
+    for _ in range(max_new_tokens):            
+        next_token = torch.argmax(out, dim=-1, keepdim=True)
+        if (eos_token_id is not None and torch.all(next_token == eos_token_id)):
+            break
+        yield next_token
+        if kv_cache:
+            out = model(next_token, cache=cache)[:, -1]
+        else:
+            token_ids = torch.cat([token_ids, next_token], dim=1)
+            out = model(token_ids)[:, -1]
+            
+def generate_text_stream_concat( model, tokenizer, prompt, device, max_new_tokens, verbose=False ):
+    input_ids = torch.tensor( tokenizer.encode(prompt), device=device ).unsqueeze(0)
+    generated_ids = []
+    for token in generate_text_basic_stream( model=model, token_ids=input_ids, max_new_tokens=max_new_tokens, eos_token_id=tokenizer.eos_token_id ):
+        next_token_id = token.squeeze(0)
+        generated_ids.append(next_token_id.item())
+        if verbose:
+            print( tokenizer.decode(next_token_id.tolist()), end="", flush=True )
+    return tokenizer.decode(generated_ids)
+
+        
+def generate_text_stream_concat_flex( model, tokenizer, prompt, device, max_new_tokens, verbose=False, generate_func=None, **generate_kwargs ):
+    if generate_func is None:
+        generate_func = generate_text_basic_stream_cache      
+    input_ids = torch.tensor( tokenizer.encode(prompt), device=device ).unsqueeze(0)
+    generated_ids = []
+    for token in generate_func(
+        model=model,
+        token_ids=input_ids,
+        max_new_tokens=max_new_tokens,
+        eos_token_id=tokenizer.eos_token_id,
+        **generate_kwargs,
+    ):
+        next_token_id = token.squeeze(0)
+        generated_ids.append(next_token_id.item())
+        if verbose:
+            print( tokenizer.decode(next_token_id.tolist()), end="", flush=True )
+    return tokenizer.decode(generated_ids)
+
+@torch.inference_mode()
+def generate_text_basic_stream_cache( model, token_ids, max_new_tokens, eos_token_id=None ):
+    model.eval()
+    cache = KVCache(n_layers=model.cfg["n_layers"])
+    model.reset_kv_cache()
+    out = model(token_ids, cache=cache)[:, -1]
+    for _ in range(max_new_tokens):
+        next_token = torch.argmax(out, dim=-1, keepdim=True)
+        if (eos_token_id is not None and torch.all(next_token == eos_token_id)):
+            break
+        yield next_token
+        out = model(next_token, cache=cache)[:, -1]
+        
+@torch.inference_mode()
+def generate_text_top_p_stream_cache( model, token_ids, max_new_tokens, eos_token_id=None, temperature=0., top_p=None ):
+    model.eval()
+    cache = KVCache(n_layers=model.cfg["n_layers"])
+    model.reset_kv_cache()
+    out = model(token_ids, cache=cache)[:, -1]
+    for _ in range(max_new_tokens):
+        orig_device = token_ids.device
+        if temperature <= 0:
+            raise ValueError("Temperature must be positive")
+        elif temperature is None or temperature == 0.0:
+            next_token = torch.argmax(out, dim=-1, keepdim=True)
+        else:            
+            logits = out / temperature
+            probas = torch.softmax(logits, dim=-1)
+            probas = top_p_filter(probas, top_p)
+            next_token = torch.multinomial(probas.cpu(), num_samples=1)
+            next_token = next_token.to(orig_device)
+        if (eos_token_id is not None
+                and torch.all(next_token == eos_token_id)):
+            break
+        yield next_token
+        out = model(next_token, cache=cache)[:, -1]
+        
 #%% Sanity Check: run from the parent folder of simonbb as "python -m simonbb.cis433.utils" since we used relative import
 
 if __name__ == "__main__":

simonbb-0.0.7/src/simonbb/cis438/agentic.py → simonbb-0.0.8/src/simonbb/cis438/rag.py

@@ -4,6 +4,7 @@
 import os
 import re
 from typing import Annotated, Literal, TypedDict, get_args
+from pathlib import Path
 
 from functools import partial
 
@@ -31,42 +32,39 @@ embedding_model = HuggingFaceEmbeddings(
 
 #%%
 
-def load_documents(path):
+def load_documents(path: Path):
     """Load documents from a file or directory."""    
 
+    valid_extensions = {'.pdf', '.html', '.ejs'}
+    
     docs = []
     success_cnt = 0
     page_cnt = 0
-
-    if os.path.isfile(path):
-        filename = os.path.basename(path)
-        print(filename)
-        ext = filename.lower()
-        if ext.endswith('.pdf'):
-            documents = PyPDFLoader(path).load()
-        elif ext.endswith(('.html', '.ejs')):
-            documents = UnstructuredHTMLLoader(path).load()
-        else:
-            print(f"Unsupported file type: {filename}")
-            return None, 0, 0
-        success_cnt = 1
-        page_cnt = len(documents)
-        docs.extend(documents)
-    elif os.path.isdir(path):
-        for filename in os.listdir(path):
-            file = os.path.join(path, filename)
-            if not os.path.isfile(file):
-                continue
-            print(filename)
-            ext = filename.lower()
-            if ext.endswith('.pdf'):
-                documents = PyPDFLoader(file).load()
-            elif ext.endswith(('.html', '.ejs')):
-                documents = UnstructuredHTMLLoader(file).load()
+    
+    if path.exists():
+        if path.is_dir():
+            for file in [ item for item in path.rglob('*') if item.is_file() and item.suffix in valid_extensions ]:
+                print(file.name)                
+                if file.suffix == '.pdf':
+                    documents = PyPDFLoader(file).load()
+                elif file.suffix in {'.html', '.ejs'}:
+                    documents = UnstructuredHTMLLoader(file).load()
+                else:
+                    continue
+                success_cnt += 1
+                page_cnt += len(documents)
+                docs.extend(documents)
+        else:            
+            print(path.name)
+            if path.suffix == '.pdf':
+                documents = PyPDFLoader(path).load()
+            elif path.suffix in {'.html', '.ejs'}:
+                documents = UnstructuredHTMLLoader(path).load()
             else:
-                continue
-            success_cnt += 1
-            page_cnt += len(documents)
+                print(f"Unsupported file type: {path}")
+                return None, 0, 0
+            success_cnt = 1
+            page_cnt = len(documents)
             docs.extend(documents)
     else:
         print(f"Error: '{path}' is not a valid file or directory")

{simonbb-0.0.7 → simonbb-0.0.8}/src/simonbb/utils.py

@@ -1,15 +1,18 @@
 import time
 import re
+from IPython.display import display, Math
 import collections
 import sys
 import pickle
 import json
 from collections import Counter
 import urllib.request
+
 import numpy as np
 import torch
 import torch.nn as nn
 
+
 class WeightedMSELoss(nn.Module):
     def __init__(self, size_average=True):
         super(WeightedMSELoss, self).__init__()
@@ -268,3 +271,213 @@ def query_ollama(prompt, model="llama3", url="http://localhost:11434/api/chat"):
             response_data += response_json["message"]["content"]
 
     return response_data
+
+def display_latex( prose_and_equations ):
+    pattern = re.compile( r"\$\$(.*?)\$\$|(?<!\$)\$(?!\$)(.*?)(?<!\$)\$(?!\$)", flags=re.DOTALL )
+    position = 0
+    for match in pattern.finditer( prose_and_equations ):        
+        print( prose_and_equations[position:match.start()], end="") # Preserve and print all prose exactly
+        equation = match.group(1) if match.group(1) is not None else match.group(2)
+        display(Math(equation.strip()))
+        position = match.end()    
+    print(prose_and_equations[position:], end="") # Preserve text following the final equation
+    
+def eta_progress_message( processed, total, start_time, show_eta=False, label="Progress" ):
+    progress = f"{label}: {processed}/{total}"
+    pad_width = len(f"{label}: {total}/{total} | ETA: 00h 00m 00s")
+    if not show_eta or processed <= 0:
+        return progress.ljust(pad_width)
+    elapsed = time.time() - start_time
+    if elapsed <= 0:
+        return progress.ljust(pad_width)
+    remaining = max(total - processed, 0)
+    if processed:
+        avg_time = elapsed / processed
+        eta_seconds = avg_time * remaining
+    else:
+        eta_seconds = 0
+    eta_seconds = max(int(round(eta_seconds)), 0)
+    minutes, rem_seconds = divmod(eta_seconds, 60)
+    hours, minutes = divmod(minutes, 60)
+    if hours:
+        eta = f"{hours}h {minutes:02d}m {rem_seconds:02d}s"
+    elif minutes:
+        eta = f"{minutes:02d}m {rem_seconds:02d}s"
+    else:
+        eta = f"{rem_seconds:02d}s"
+    message = f"{progress} | ETA: {eta}"
+    return message.ljust(pad_width)
+
+#%% math verifier
+
+from sympy import simplify
+from sympy.parsing import sympy_parser
+from sympy.core.sympify import SympifyError
+from tokenize import TokenError
+from sympy.polys.polyerrors import PolynomialError
+
+def grade_math_answer(pred_text, gt_text):
+    result = False  # Default outcome if checks fail
+    # Only continue if both inputs are non-empty strings
+    if pred_text is not None and gt_text is not None:
+        gt_parts = _split_into_parts( _normalize_text(gt_text) )  # Break ground truth into comparable parts
+        pred_parts = _split_into_parts( _normalize_text(pred_text) )  # Break prediction into comparable parts
+        # Ensure both sides have same number of valid parts
+        if (gt_parts and pred_parts
+           and len(gt_parts) == len(pred_parts)):
+            result = all(
+                equality_check(gt, pred)
+                for gt, pred in zip(gt_parts, pred_parts)
+            )  # Check each part for mathematical equivalence
+    return result  # True only if all checks passed
+
+
+def equality_check(expr_gtruth, expr_pred):
+    # First, check if the two expressions are exactly the same string
+    if expr_gtruth == expr_pred:
+        return True
+    # Parse both expressions into SymPy objects (returns None if parsing fails)
+    gtruth, pred = _sympy_parser(expr_gtruth), _sympy_parser(expr_pred)
+    # If both expressions were parsed successfully, try symbolic comparison
+    if gtruth is not None and pred is not None:
+        try:
+            # If the difference is 0, they are equivalent
+            return simplify(gtruth - pred) == 0
+        except (SympifyError, TypeError):
+            pass
+    return False
+
+def _split_into_parts(text):
+    result = [text]
+    if text:
+        # Check if text looks like a tuple or list, e.g. "(a, b)" or "[a, b]"
+        if (
+            len(text) >= 2
+            and text[0] in "([" and text[-1] in ")]"
+            and "," in text[1:-1]
+        ):
+            # Split on commas inside brackets and strip whitespace
+            items = [p.strip() for p in text[1:-1].split(",")]
+            if all(items):
+                result = items
+    else:
+        # If text is empty, return an empty list
+        result = []
+
+    return result
+
+
+def _sympy_parser(expr):
+    # To avoid crashing on long garbage responses 
+    # that some badly trained models (chapter 6) may emit
+    if expr is None or len(expr) > 2000:
+        return None
+    try:
+        return sympy_parser.parse_expr(
+            expr,
+            transformations=(
+                # Standard transformations like handling parentheses
+                *sympy_parser.standard_transformations,
+                # Allow omitted multiplication symbols (e.g., "2x" -> 2*x")
+                sympy_parser.implicit_multiplication_application,
+            ),
+            # Evaluate during parsing so simple constants simplify (e.g., 2+3 -> 5)
+            evaluate=True,
+        )
+    except (SympifyError, SyntaxError, TypeError, AttributeError,
+            IndexError, TokenError, ValueError, PolynomialError):
+        return None
+    
+LATEX_FIXES = [  # Latex formatting to be replaced
+    (r"\\left\s*", ""),
+    (r"\\right\s*", ""),
+    (r"\\,|\\!|\\;|\\:", ""),
+    (r"\\cdot", "*"),
+    (r"\u00B7|\u00D7", "*"),
+    (r"\\\^\\circ", ""),
+    (r"\\dfrac", r"\\frac"),
+    (r"\\tfrac", r"\\frac"),
+    (r"°", ""),
+]
+
+RE_SPECIAL = re.compile(r"<\|[^>]+?\|>")  # strip chat special tokens like <|assistant|>
+
+def _normalize_text(text):
+    if not text:
+        return ""
+    text = RE_SPECIAL.sub("", text).strip()
+    SUPERSCRIPT_MAP = {
+        "⁰": "0", "¹": "1", "²": "2", "³": "3", "⁴": "4",
+        "⁵": "5", "⁶": "6", "⁷": "7", "⁸": "8", "⁹": "9",
+        "⁺": "+", "⁻": "-", "⁽": "(", "⁾": ")",
+    }
+    # Strip leading multiple-choice labels
+    # E.g., like "c. 3" -> 3, or "b: 2" -> 2
+    match = re.match(r"^[A-Za-z]\s*[.:]\s*(.+)$", text)
+    if match:
+        text = match.group(1)        
+    # Remove angle-degree markers
+    text = re.sub(r"\^\s*\{\s*\\circ\s*\}", "", text)   # ^{\circ}
+    text = re.sub(r"\^\s*\\circ", "", text)             # ^\circ
+    text = text.replace("°", "")                        # Unicode degree
+    # unwrap \text{...} if the whole string is wrapped
+    match = re.match(r"^\\text\{(?P<x>.+?)\}$", text)
+    if match:
+        text = match.group("x")
+    # strip inline/display math wrappers \( \) \[ \]
+    text = re.sub(r"\\\(|\\\)|\\\[|\\\]", "", text)
+    # light LaTeX canonicalization
+    for pat, rep in LATEX_FIXES:
+        text = re.sub(pat, rep, text)
+    def convert_superscripts(s, base=None):
+        converted = "".join(
+            SUPERSCRIPT_MAP[ch] if ch in SUPERSCRIPT_MAP else ch
+            for ch in s
+        )
+        if base is None:
+            return converted
+        return f"{base}**{converted}"
+    # convert unicode superscripts into exponent form (e.g., 2² -> 2**2)m
+    text = re.sub(
+        r"([0-9A-Za-z\)\]\}])([⁰¹²³⁴⁵⁶⁷⁸⁹⁺⁻]+)",
+        lambda m: convert_superscripts(m.group(2), base=m.group(1)),
+        text,
+    )
+    text = convert_superscripts(text)    
+    # numbers/roots
+    text = text.replace("\\%", "%").replace("$", "").replace("%", "")
+    text = re.sub(
+        r"\\sqrt\s*\{([^}]*)\}",
+        lambda match: f"sqrt({match.group(1)})",
+        text,
+    )
+    text = re.sub(
+        r"\\sqrt\s+([^\\\s{}]+)",
+        lambda match: f"sqrt({match.group(1)})",
+        text,
+    )
+    # fractions
+    text = re.sub(
+        r"\\frac\s*\{([^{}]+)\}\s*\{([^{}]+)\}",
+        lambda match: f"({match.group(1)})/({match.group(2)})",
+        text,
+    )
+    text = re.sub(
+        r"\\frac\s+([^\s{}]+)\s+([^\s{}]+)",
+        lambda match: f"({match.group(1)})/({match.group(2)})",
+        text,
+    )
+    # exponent and mixed numbers
+    text = text.replace("^", "**")
+    text = re.sub(
+        r"(?<=\d)\s+(\d+/\d+)",
+        lambda match: "+" + match.group(1),
+        text,
+    )
+    # 1,234 -> 1234
+    text = re.sub(
+        r"(?<=\d),(?=\d\d\d(\D|$))",
+        "",
+        text,
+    )
+    return text.replace("{", "").replace("}", "").strip().lower()