| import argparse |
| import json |
| import sys |
| from pathlib import Path |
| import re |
|
|
| import torch |
| import torch.nn.functional as F |
| from transformers import AutoTokenizer |
|
|
| try: |
| from tokenizers import Tokenizer as HFTokenizer |
| except ImportError: |
| HFTokenizer = None |
|
|
| DEFAULT_CHECKPOINT = r"D:\Downloads\hssm_fineweb_edu_final.pt" |
| DEFAULT_TOKENIZER = r"D:\Downloads\simple_tokenizer_20k.json" |
| RUBINET_HSSM_PATH = r"C:\Users\ASUS\.anaconda" |
|
|
| sys.path.append(RUBINET_HSSM_PATH) |
| from RubiNet_HSSM import HierarchicalSSM |
| from hssm_v2_gpu_pretrain import HSSMV2Config, HSSMV2LM |
|
|
| if hasattr(sys.stdout, "reconfigure"): |
| sys.stdout.reconfigure(encoding="utf-8", errors="replace") |
| if hasattr(sys.stderr, "reconfigure"): |
| sys.stderr.reconfigure(encoding="utf-8", errors="replace") |
|
|
|
|
| class CompatibleTokenizer: |
| def __init__(self, tokenizer_path: str): |
| path = Path(tokenizer_path) |
| with path.open("r", encoding="utf-8") as f: |
| data = json.load(f) |
|
|
| self.backend = None |
| if HFTokenizer is not None: |
| try: |
| self.backend = HFTokenizer.from_file(str(path)) |
| except Exception: |
| self.backend = None |
|
|
| if "model" in data and isinstance(data["model"], dict) and "vocab" in data["model"]: |
| vocab = data["model"]["vocab"] |
| elif "vocab" in data: |
| vocab = data["vocab"] |
| else: |
| raise ValueError(f"Unsupported tokenizer format: {path}") |
|
|
| self.vocab = {str(token): int(idx) for token, idx in vocab.items()} |
| self.id_to_token = {idx: token for token, idx in self.vocab.items()} |
| self.vocab_size = len(self.vocab) |
| self.pad_token_id = self._resolve_token_id(["<PAD>", "[PAD]"], fallback=0) |
| self.unk_token_id = self._resolve_token_id(["<UNK>", "[UNK]"], fallback=3) |
|
|
| print(f"[TOKENIZER] Loaded vocab tokenizer - Vocab: {self.vocab_size:,}") |
|
|
| def _resolve_token_id(self, candidates, fallback: int): |
| for token in candidates: |
| token_id = self.vocab.get(token) |
| if token_id is not None: |
| return token_id |
| return fallback |
|
|
| def encode(self, text, max_length=128): |
| if self.backend is not None: |
| ids = self.backend.encode(text).ids[:max_length] |
| else: |
| words = text.split() |
| ids = [self.vocab.get(word, self.unk_token_id) for word in words][:max_length] |
| if len(ids) < max_length: |
| ids += [self.pad_token_id] * (max_length - len(ids)) |
| return ids |
|
|
| def decode(self, ids): |
| filtered = [int(i) for i in ids if int(i) != self.pad_token_id] |
| if self.backend is not None: |
| return self.backend.decode(filtered, skip_special_tokens=False) |
| return " ".join(self.id_to_token.get(i, "<UNK>") for i in filtered) |
|
|
|
|
| class HFTokenizerAdapter: |
| def __init__(self, tokenizer_name: str): |
| self.backend = AutoTokenizer.from_pretrained(tokenizer_name, use_fast=True) |
| if self.backend.pad_token is None: |
| self.backend.pad_token = self.backend.eos_token or self.backend.unk_token |
| self.vocab = self.backend.get_vocab() |
| self.id_to_token = {idx: token for token, idx in self.vocab.items()} |
| self.vocab_size = int(self.backend.vocab_size) |
| self.pad_token_id = int(self.backend.pad_token_id) |
| self.unk_token_id = int(self.backend.unk_token_id if self.backend.unk_token_id is not None else self.pad_token_id) |
|
|
| def encode(self, text, max_length=128): |
| ids = self.backend.encode(text, add_special_tokens=False, truncation=True, max_length=max_length) |
| if len(ids) < max_length: |
| ids += [self.pad_token_id] * (max_length - len(ids)) |
| return ids |
|
|
| def decode(self, ids): |
| filtered = [int(i) for i in ids if int(i) != self.pad_token_id] |
| return self.backend.decode(filtered, skip_special_tokens=True) |
|
|
|
|
| def build_model(tokenizer): |
| return HierarchicalSSM( |
| vocab_size=tokenizer.vocab_size, |
| d_model=512, |
| d_state=32, |
| num_blocks=6, |
| num_experts=8, |
| top_k=2, |
| chunk_size=4, |
| expert_dim=1024, |
| ) |
|
|
|
|
| def build_hssm_v2_model(tokenizer, checkpoint_config: dict): |
| config = HSSMV2Config( |
| vocab_size=int(checkpoint_config.get("vocab_size", tokenizer.vocab_size)), |
| d_model=int(checkpoint_config.get("d_model", 288)), |
| n_layers=int(checkpoint_config.get("n_layers", 10)), |
| d_ff=int(checkpoint_config.get("d_ff", 512)), |
| state_rank=int(checkpoint_config.get("state_rank", 128)), |
| chunk_size=int(checkpoint_config.get("chunk_size", 8)), |
| dropout=float(checkpoint_config.get("dropout", 0.0)), |
| max_seq_len=int(checkpoint_config.get("max_seq_len", 1024)), |
| tie_embeddings=bool(checkpoint_config.get("tie_embeddings", True)), |
| num_experts=int(checkpoint_config.get("num_experts", 64)), |
| experts_per_token=int(checkpoint_config.get("experts_per_token", 1)), |
| expert_dim=int(checkpoint_config.get("expert_dim", 2048)), |
| moe_every=int(checkpoint_config.get("moe_every", 4)), |
| aux_loss_coef=float(checkpoint_config.get("aux_loss_coef", 1e-2)), |
| ) |
| return HSSMV2LM(config) |
|
|
|
|
| def _looks_like_hf_tokenizer_reference(tokenizer_path: str) -> bool: |
| path = Path(tokenizer_path) |
| return not path.exists() |
|
|
|
|
| def _load_tokenizer(tokenizer_path: str): |
| if _looks_like_hf_tokenizer_reference(tokenizer_path): |
| return HFTokenizerAdapter(tokenizer_path) |
| return CompatibleTokenizer(tokenizer_path) |
|
|
|
|
| def _is_hssm_v2_checkpoint(checkpoint: dict) -> bool: |
| config = checkpoint.get("config") if isinstance(checkpoint, dict) else None |
| if not isinstance(config, dict): |
| return False |
| required_keys = {"d_model", "n_layers", "state_rank", "chunk_size"} |
| return required_keys.issubset(config.keys()) |
|
|
|
|
| def load_pretrained(checkpoint_path: str, tokenizer_path: str, device: str): |
| checkpoint_file = Path(checkpoint_path) |
|
|
| if not checkpoint_file.exists(): |
| raise FileNotFoundError(f"Checkpoint not found: {checkpoint_file}") |
|
|
| if not _looks_like_hf_tokenizer_reference(tokenizer_path): |
| tokenizer_file = Path(tokenizer_path) |
| if not tokenizer_file.exists(): |
| raise FileNotFoundError(f"Tokenizer not found: {tokenizer_file}") |
|
|
| tokenizer = _load_tokenizer(tokenizer_path) |
|
|
| checkpoint = torch.load(str(checkpoint_file), map_location=device, weights_only=False) |
| state_dict = checkpoint["model_state_dict"] if "model_state_dict" in checkpoint else checkpoint |
| if _is_hssm_v2_checkpoint(checkpoint): |
| model = build_hssm_v2_model(tokenizer, checkpoint.get("config", {})) |
| else: |
| model = build_model(tokenizer) |
| missing, unexpected = model.load_state_dict(state_dict, strict=False) |
|
|
| model = model.to(device) |
| model.eval() |
|
|
| print("Loaded HSSM checkpoint") |
| print(f" Path: {checkpoint_file}") |
| print(f" Missing keys: {len(missing)}") |
| print(f" Unexpected keys: {len(unexpected)}") |
| if "epoch" in checkpoint: |
| print(f" Epoch: {checkpoint['epoch']}") |
| if "loss" in checkpoint: |
| print(f" Loss: {checkpoint['loss']}") |
| print(f" Model type: {'HSSM v2' if _is_hssm_v2_checkpoint(checkpoint) else 'RubiNet HSSM'}") |
|
|
| return tokenizer, model |
|
|
|
|
| def _model_chunk_size(model) -> int: |
| if hasattr(model, "chunk_size"): |
| return int(model.chunk_size) |
| if hasattr(model, "config") and hasattr(model.config, "chunk_size"): |
| return int(model.config.chunk_size) |
| return 1 |
|
|
|
|
| def _next_token_logits(model, input_tensor: torch.Tensor, current_len: int) -> torch.Tensor: |
| outputs = model(input_tensor) |
| if isinstance(outputs, dict): |
| logits = outputs.get("logits") |
| if logits is None: |
| raise ValueError("Model returned a dict without logits") |
| return logits[0, current_len - 1, :].clone() |
| chunk_size = _model_chunk_size(model) |
| chunk_idx = max((current_len - 1) // chunk_size, 0) |
| return outputs[0, chunk_idx, :].clone() |
|
|
|
|
| def build_prompt(user_text: str, cot_mode: bool = False) -> str: |
| cleaned_user_text = user_text.strip() |
| if cot_mode: |
| return ( |
| "system: Reply only in correct English. " |
| "Follow English grammar, spelling, punctuation, and sentence structure strictly. " |
| "Do not output fragments, corrupted tokens, mixed-language text, or placeholder symbols. " |
| "Think step by step briefly and keep the output clean. " |
| "Output exactly two lines in this format: " |
| "Reasoning: <very short reasoning>. " |
| "Answer: <final answer>. " |
| "Keep both lines grammatical and concise.\n" |
| f"user: {cleaned_user_text}\n" |
| "assistant:" |
| ) |
| return ( |
| "system: Reply only in correct English. " |
| "Follow English grammar, spelling, punctuation, and sentence structure strictly. " |
| "Use short complete sentences. " |
| "Do not output broken words, malformed tokens, mixed-language text, or placeholder symbols.\n" |
| f"user: {cleaned_user_text}\n" |
| "assistant:" |
| ) |
|
|
|
|
| def safe_print(text: str): |
| try: |
| print(text) |
| except UnicodeEncodeError: |
| sanitized = text.encode("utf-8", errors="replace").decode("utf-8", errors="replace") |
| print(sanitized) |
|
|
|
|
| def _apply_top_p_filter(logits: torch.Tensor, top_p: float) -> torch.Tensor: |
| if top_p >= 1.0: |
| return logits |
| sorted_logits, sorted_indices = torch.sort(logits, descending=True) |
| sorted_probs = F.softmax(sorted_logits, dim=-1) |
| cumulative_probs = torch.cumsum(sorted_probs, dim=-1) |
| sorted_indices_to_remove = cumulative_probs > top_p |
| sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() |
| sorted_indices_to_remove[..., 0] = False |
| indices_to_remove = sorted_indices[sorted_indices_to_remove] |
| logits[indices_to_remove] = float("-inf") |
| return logits |
|
|
|
|
| def _has_repeat_ngram(token_ids, next_token_id: int, ngram_size: int) -> bool: |
| if ngram_size <= 1 or len(token_ids) < ngram_size - 1: |
| return False |
| candidate = token_ids[-(ngram_size - 1):] + [next_token_id] |
| for i in range(len(token_ids) - ngram_size + 1): |
| if token_ids[i:i + ngram_size] == candidate: |
| return True |
| return False |
|
|
|
|
| def _normalize_word(word: str) -> str: |
| return re.sub(r"[^a-z0-9]+", "", word.lower()) |
|
|
|
|
| def _recent_word_counts(text: str, window: int = 12): |
| words = [_normalize_word(part) for part in text.split()] |
| words = [word for word in words if word] |
| recent = words[-window:] |
| counts = {} |
| for word in recent: |
| counts[word] = counts.get(word, 0) + 1 |
| return counts |
|
|
|
|
| def _violates_word_repeat(decoded_text: str, candidate_piece: str) -> bool: |
| candidate_word = _normalize_word(candidate_piece) |
| if not candidate_word: |
| return False |
| counts = _recent_word_counts(decoded_text, window=12) |
| return counts.get(candidate_word, 0) >= 2 |
|
|
|
|
| def _resolve_special_token_ids(tokenizer): |
| special_ids = set() |
| for token in ["<BOS>", "[BOS]", "<PAD>", "[PAD]", "<SEP>", "[SEP]", "<EOS>", "[EOS]", "<UNK>", "[UNK]", "<CLS>", "[CLS]", "<MASK>", "[MASK]", "<MASK>"]: |
| token_id = tokenizer.vocab.get(token) |
| if token_id is not None: |
| special_ids.add(int(token_id)) |
| if getattr(tokenizer, "pad_token_id", None) is not None: |
| special_ids.add(int(tokenizer.pad_token_id)) |
| if getattr(tokenizer, "unk_token_id", None) is not None: |
| special_ids.add(int(tokenizer.unk_token_id)) |
| return special_ids |
|
|
|
|
| def _contains_special_marker(text: str) -> bool: |
| upper_text = text.upper() |
| markers = ["<BOS>", "[BOS]", "<PAD>", "[PAD]", "<SEP>", "[SEP]", "<EOS>", "[EOS]", "<UNK>", "[UNK]", "<CLS>", "[CLS]", "<MASK>", "[MASK]"] |
| return any(marker in upper_text for marker in markers) |
|
|
|
|
| def _looks_like_artifact(text: str) -> bool: |
| stripped = text.strip() |
| if not stripped: |
| return False |
| if "##" in stripped: |
| return True |
| if stripped.startswith("##") or stripped.endswith("##"): |
| return True |
| if stripped.count("#") >= 1 and len(stripped) <= 4: |
| return True |
| if "�" in stripped: |
| return True |
| if any(ch in stripped for ch in ["�", ""]): |
| return True |
| if re.search(r"(.)\1{3,}", stripped.lower()): |
| return True |
| if re.fullmatch(r"[A-Za-z]{1,4}\d{2,}", stripped): |
| return True |
| if re.fullmatch(r"[#\-_=~`|.]+", stripped): |
| return True |
| return False |
|
|
|
|
| def _strip_special_markers(text: str) -> str: |
| cleaned = text |
| for pattern in [r"<\s*BOS\s*>", r"\[\s*BOS\s*\]", r"<\s*PAD\s*>", r"\[\s*PAD\s*\]", r"<\s*SEP\s*>", r"\[\s*SEP\s*\]", r"<\s*EOS\s*>", r"\[\s*EOS\s*\]", r"<\s*UNK\s*>", r"\[\s*UNK\s*\]", r"<\s*CLS\s*>", r"\[\s*CLS\s*\]", r"<\s*MASK\s*>", r"\[\s*MASK\s*\]"]: |
| cleaned = re.sub(pattern, " ", cleaned, flags=re.IGNORECASE) |
| cleaned = re.sub(r"#{2,}", " ", cleaned) |
| cleaned = re.sub(r"(?<!\w)#(?!\w)", " ", cleaned) |
| cleaned = cleaned.replace("�", " ") |
| cleaned = re.sub(r"\b([A-Za-z]+)(\s+\1\b){2,}", r"\1", cleaned, flags=re.IGNORECASE) |
| cleaned = re.sub(r"\b(\w{1,20})(\w{1,20})\1\b", r"\1", cleaned) |
| cleaned = re.sub(r"\s*([,;:.!?])\s*", r"\1 ", cleaned) |
| cleaned = re.sub(r"\s+", " ", cleaned) |
| return cleaned.strip() |
|
|
|
|
| def _cleanup_english_grammar(text: str) -> str: |
| cleaned = text.strip() |
| if not cleaned: |
| return cleaned |
|
|
| replacements = { |
| " im ": " I'm ", |
| " ive ": " I've ", |
| " ill ": " I'll ", |
| " id ": " I'd ", |
| " dont ": " don't ", |
| " cant ": " can't ", |
| " wont ": " won't ", |
| " didnt ": " didn't ", |
| " doesnt ": " doesn't ", |
| " isnt ": " isn't ", |
| " arent ": " aren't ", |
| " wasnt ": " wasn't ", |
| " werent ": " weren't ", |
| " thats ": " that's ", |
| " whats ": " what's ", |
| " theres ": " there's ", |
| " ive ": " I've ", |
| } |
|
|
| padded = f" {cleaned} " |
| for source, target in replacements.items(): |
| padded = re.sub(re.escape(source), target, padded, flags=re.IGNORECASE) |
| cleaned = padded.strip() |
|
|
| cleaned = re.sub(r"\bi\b", "I", cleaned) |
| cleaned = re.sub(r"\b([A-Za-z]+)(\s+\1\b){1,}", r"\1", cleaned, flags=re.IGNORECASE) |
| cleaned = re.sub(r"\s+([,;:.!?])", r"\1", cleaned) |
| cleaned = re.sub(r"([,;:.!?])(?!\s|$)", r"\1 ", cleaned) |
| cleaned = re.sub(r"\s+", " ", cleaned).strip() |
|
|
| if cleaned: |
| cleaned = cleaned[0].upper() + cleaned[1:] |
|
|
| sentences = re.split(r"(?<=[.!?])\s+", cleaned) |
| normalized_sentences = [] |
| for sentence in sentences: |
| sentence = sentence.strip() |
| if not sentence: |
| continue |
| if len(sentence) == 1: |
| normalized_sentences.append(sentence.upper()) |
| else: |
| normalized_sentences.append(sentence[0].upper() + sentence[1:]) |
| cleaned = " ".join(normalized_sentences).strip() |
|
|
| if cleaned and cleaned[-1] not in ".!?": |
| cleaned += "." |
|
|
| return cleaned |
|
|
|
|
| def _is_strict_english_output(text: str, cot_mode: bool = False) -> bool: |
| cleaned = text.strip() |
| if not cleaned: |
| return False |
| if any(token in cleaned for token in ["[", "]", "{", "}", "|", "<UNK>", "[UNK]", "<PAD>", "[PAD]"]): |
| return False |
| if re.search(r"[^A-Za-z0-9\s,.;:!?\-\'\"()\n]", cleaned): |
| return False |
| words = re.findall(r"[A-Za-z']+", cleaned) |
| if len(words) < 2: |
| return False |
| long_weird_words = [word for word in words if len(word) > 18] |
| if long_weird_words: |
| return False |
| if re.search(r"([A-Za-z]{2,})([A-Z][a-z]+)", cleaned): |
| return False |
| common_markers = { |
| "the", "a", "an", "is", "are", "am", "i", "you", "we", "they", "it", "to", "of", "and", |
| "that", "this", "can", "will", "do", "not", "yes", "no", "my", "your", "in", "on", "for" |
| } |
| lowered_words = [word.lower() for word in words] |
| if not any(word in common_markers for word in lowered_words): |
| return False |
| if cot_mode: |
| lines = [line.strip() for line in cleaned.splitlines() if line.strip()] |
| if len(lines) != 2: |
| return False |
| if not lines[0].startswith("Reasoning:"): |
| return False |
| if not lines[1].startswith("Answer:"): |
| return False |
| sentences = [segment.strip() for segment in re.split(r"(?<=[.!?])\s+", cleaned) if segment.strip()] |
| if not sentences: |
| return False |
| for sentence in sentences: |
| if not sentence[0].isupper(): |
| return False |
| if sentence[-1] not in ".!?": |
| return False |
| return True |
|
|
|
|
| def _force_cot_shape(text: str) -> str: |
| cleaned = text.strip() |
| if not cleaned: |
| return cleaned |
| lines = [line.strip() for line in cleaned.splitlines() if line.strip()] |
| if len(lines) >= 2 and lines[0].startswith("Reasoning:") and lines[1].startswith("Answer:"): |
| return f"{lines[0]}\n{lines[1]}" |
| parts = re.split(r"(?<=[.!?])\s+", cleaned, maxsplit=1) |
| if len(parts) == 2: |
| reasoning, answer = parts |
| else: |
| reasoning, answer = "Reasoning: Briefly considered the request.", f"Answer: {cleaned}" |
| return f"{reasoning}\n{answer}" |
| reasoning = reasoning if reasoning.startswith("Reasoning:") else f"Reasoning: {reasoning.strip()}" |
| answer = answer if answer.startswith("Answer:") else f"Answer: {answer.strip()}" |
| return f"{reasoning}\n{answer}" |
|
|
|
|
| def _ban_low_quality_candidates(tokenizer, logits: torch.Tensor): |
| for token_id in range(logits.size(0)): |
| piece = tokenizer.decode([token_id]).strip() |
| if not piece: |
| continue |
| if _contains_special_marker(piece): |
| logits[token_id] = float("-inf") |
|
|
|
|
| def _select_candidate_id(tokenizer, probs: torch.Tensor, generated, prompt_token_count: int, no_repeat_ngram_size: int): |
| candidate_count = min(24, probs.size(0)) |
| top_probs, top_ids = torch.topk(probs, candidate_count) |
| decoded_so_far = tokenizer.decode(generated[prompt_token_count:]).strip() |
|
|
| fallback_clean_id = None |
| fallback_clean_prob = -1.0 |
| fallback_any_id = None |
| fallback_any_prob = -1.0 |
| for prob_value, candidate_id_tensor in zip(top_probs.tolist(), top_ids.tolist()): |
| candidate_id = int(candidate_id_tensor) |
| candidate_piece = tokenizer.decode([candidate_id]).strip() |
| if not candidate_piece: |
| continue |
| if _contains_special_marker(candidate_piece): |
| continue |
| if fallback_any_id is None or prob_value > fallback_any_prob: |
| fallback_any_id = candidate_id |
| fallback_any_prob = prob_value |
| if _looks_like_artifact(candidate_piece): |
| continue |
| if _violates_word_repeat(decoded_so_far, candidate_piece): |
| continue |
| if _has_repeat_ngram(generated, candidate_id, max(no_repeat_ngram_size, 4)): |
| continue |
| normalized_piece = _normalize_word(candidate_piece) |
| if normalized_piece and decoded_so_far: |
| recent_words = [_normalize_word(part) for part in decoded_so_far.split()[-8:]] |
| recent_words = [word for word in recent_words if word] |
| if recent_words.count(normalized_piece) >= 1: |
| continue |
| if fallback_clean_id is None or prob_value > fallback_clean_prob: |
| fallback_clean_id = candidate_id |
| fallback_clean_prob = prob_value |
|
|
| if fallback_clean_id is not None: |
| return fallback_clean_id |
| return fallback_any_id |
|
|
|
|
| def _generate_fallback_reply(model, tokenizer, prompt_tokens, blocked_special_ids, max_length: int): |
| device = next(model.parameters()).device |
| generated = list(prompt_tokens) |
|
|
| with torch.no_grad(): |
| for _ in range(min(max_length, 16)): |
| current_len = len(generated) |
| chunk_size = _model_chunk_size(model) |
| pad_len = (chunk_size - current_len % chunk_size) % chunk_size |
| padded_input = generated + [tokenizer.pad_token_id] * pad_len |
| input_tensor = torch.tensor([padded_input], device=device) |
| next_token_logits = _next_token_logits(model, input_tensor, current_len) |
|
|
| for special_id in blocked_special_ids: |
| if 0 <= special_id < next_token_logits.size(0): |
| next_token_logits[special_id] = float("-inf") |
|
|
| next_token_id = int(torch.argmax(next_token_logits).item()) |
| if next_token_id == tokenizer.pad_token_id: |
| break |
|
|
| next_piece = tokenizer.decode([next_token_id]).strip() |
| if not next_piece or _contains_special_marker(next_piece): |
| break |
|
|
| generated.append(next_token_id) |
|
|
| return generated |
|
|
|
|
| def generate_reply( |
| model, |
| tokenizer, |
| prompt: str, |
| max_length: int, |
| temperature: float, |
| top_k: int, |
| top_p: float, |
| repetition_penalty: float, |
| no_repeat_ngram_size: int, |
| cot_mode: bool = False, |
| ): |
| model.eval() |
| device = next(model.parameters()).device |
|
|
| formatted_prompt = build_prompt(prompt, cot_mode=cot_mode) |
| prompt_ids = tokenizer.encode(formatted_prompt, max_length=128) |
| generated = [tok for tok in prompt_ids if tok != tokenizer.pad_token_id] |
| prompt_token_count = len(generated) |
| blocked_special_ids = _resolve_special_token_ids(tokenizer) |
| bos_token_id = tokenizer.vocab.get("<BOS>") |
| if not generated: |
| generated = [tokenizer.unk_token_id] |
| prompt_token_count = len(generated) |
|
|
| with torch.no_grad(): |
| for _ in range(max_length): |
| current_len = len(generated) |
| chunk_size = _model_chunk_size(model) |
| pad_len = (chunk_size - current_len % chunk_size) % chunk_size |
| padded_input = generated + [tokenizer.pad_token_id] * pad_len |
| input_tensor = torch.tensor([padded_input], device=device) |
|
|
| next_token_logits = _next_token_logits(model, input_tensor, current_len) |
|
|
| if temperature > 0: |
| next_token_logits = next_token_logits / temperature |
|
|
| for special_id in blocked_special_ids: |
| if 0 <= special_id < next_token_logits.size(0): |
| next_token_logits[special_id] = float("-inf") |
|
|
| if bos_token_id is not None and 0 <= int(bos_token_id) < next_token_logits.size(0): |
| next_token_logits[int(bos_token_id)] = float("-inf") |
|
|
| _ban_low_quality_candidates(tokenizer, next_token_logits) |
|
|
| recent_tokens = generated[-48:] |
| recent_weights = {} |
| for idx, token_id in enumerate(recent_tokens): |
| distance_weight = 1.0 + (idx / max(len(recent_tokens), 1)) |
| recent_weights[token_id] = max(recent_weights.get(token_id, 1.0), distance_weight) |
|
|
| for token_id, distance_weight in recent_weights.items(): |
| if 0 <= token_id < next_token_logits.size(0): |
| penalty = repetition_penalty * distance_weight |
| if next_token_logits[token_id] > 0: |
| next_token_logits[token_id] /= penalty |
| else: |
| next_token_logits[token_id] *= penalty |
|
|
| for token_id in range(next_token_logits.size(0)): |
| if _has_repeat_ngram(generated, token_id, no_repeat_ngram_size): |
| next_token_logits[token_id] = float("-inf") |
|
|
| if top_k > 0 and top_k < next_token_logits.size(0): |
| threshold = torch.topk(next_token_logits, top_k)[0][..., -1] |
| next_token_logits[next_token_logits < threshold] = float("-inf") |
|
|
| next_token_logits = _apply_top_p_filter(next_token_logits, top_p) |
| probs = F.softmax(next_token_logits, dim=-1) |
|
|
| if torch.isnan(probs).any() or torch.isinf(probs).any() or probs.sum() <= 0: |
| break |
|
|
| next_token = _select_candidate_id( |
| tokenizer, |
| probs, |
| generated, |
| prompt_token_count, |
| no_repeat_ngram_size, |
| ) |
|
|
| if next_token is None: |
| break |
|
|
| if next_token == tokenizer.pad_token_id: |
| break |
| generated.append(next_token) |
|
|
| decoded_output = tokenizer.decode(generated[prompt_token_count:]).strip() |
| if len(decoded_output.split()) >= 6: |
| tail_words = [_normalize_word(part) for part in decoded_output.split()[-4:]] |
| tail_words = [word for word in tail_words if word] |
| if len(tail_words) >= 4 and len(set(tail_words)) == 1: |
| break |
|
|
| output_ids = generated[prompt_token_count:] |
| cleaned_output = _strip_special_markers(tokenizer.decode(output_ids).strip()) |
| if cleaned_output: |
| normalized_output = _cleanup_english_grammar(cleaned_output) |
| if cot_mode: |
| normalized_output = _force_cot_shape(normalized_output) |
| return normalized_output |
|
|
| fallback_generated = _generate_fallback_reply( |
| model, |
| tokenizer, |
| generated[:prompt_token_count], |
| blocked_special_ids, |
| max_length, |
| ) |
| fallback_output_ids = fallback_generated[prompt_token_count:] |
| fallback_output = _strip_special_markers(tokenizer.decode(fallback_output_ids).strip()) |
| normalized_fallback_output = _cleanup_english_grammar(fallback_output) |
| if cot_mode: |
| normalized_fallback_output = _force_cot_shape(normalized_fallback_output) |
| return normalized_fallback_output |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Chat/test with pretrained RubiNet HSSM checkpoint") |
| parser.add_argument("--checkpoint", default=DEFAULT_CHECKPOINT) |
| parser.add_argument("--tokenizer", default=DEFAULT_TOKENIZER) |
| parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") |
| parser.add_argument("--max-length", type=int, default=40) |
| parser.add_argument("--temperature", type=float, default=0.0) |
| parser.add_argument("--top-k", type=int, default=4) |
| parser.add_argument("--top-p", type=float, default=0.65) |
| parser.add_argument("--repetition-penalty", type=float, default=1.9) |
| parser.add_argument("--no-repeat-ngram-size", type=int, default=6) |
| parser.add_argument("--cot-mode", action="store_true") |
| parser.add_argument("--no-cot-mode", action="store_false", dest="cot_mode") |
| parser.set_defaults(cot_mode=False) |
| parser.add_argument("--message", default="") |
| args = parser.parse_args() |
|
|
| tokenizer, model = load_pretrained(args.checkpoint, args.tokenizer, args.device) |
|
|
| if args.message: |
| output = generate_reply( |
| model, |
| tokenizer, |
| args.message, |
| args.max_length, |
| args.temperature, |
| args.top_k, |
| args.top_p, |
| args.repetition_penalty, |
| args.no_repeat_ngram_size, |
| args.cot_mode, |
| ) |
| safe_print(output) |
| return |
|
|
| print("Interactive HSSM chat/test. Type 'exit' to quit.") |
| while True: |
| user_text = input("You: ").strip() |
| if not user_text: |
| continue |
| if user_text.lower() in {"exit", "quit"}: |
| break |
| output = generate_reply( |
| model, |
| tokenizer, |
| user_text, |
| args.max_length, |
| args.temperature, |
| args.top_k, |
| args.top_p, |
| args.repetition_penalty, |
| args.no_repeat_ngram_size, |
| args.cot_mode, |
| ) |
| safe_print(f"HSSM: {output}\n") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
