Files
soprano-factory/analyze_tokens.py
T
2026-02-09 01:01:44 -05:00

54 lines
1.8 KiB
Python

# analyze_tokens.py - Diagnostic for LLM Audio Tokens
import torch
import argparse
from transformers import AutoTokenizer, AutoModelForCausalLM
from utils.text_normalizer import normalize_text
from utils.config import cfg
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--text", type=str, required=True)
parser.add_argument("--model-dir", type=str, required=True)
args = parser.parse_args()
tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
llm = AutoModelForCausalLM.from_pretrained(args.model_dir, torch_dtype=torch.float32).to(cfg.device)
norm_text = normalize_text(args.text)
prompt = f"[STOP][TEXT]{norm_text}[START]"
inputs = tokenizer(prompt, return_tensors="pt").to(cfg.device)
print(f"Analyzing LLM response for: {norm_text}")
with torch.no_grad():
output_ids = llm.generate(
inputs["input_ids"],
max_new_tokens=128,
do_sample=True,
top_k=50
)
# Convert to list and find the start of audio
tokens = output_ids[0].tolist()
start_token_id = tokenizer.convert_tokens_to_ids("[START]")
try:
idx = tokens.index(start_token_id)
audio_tokens = tokens[idx+1:]
except:
audio_tokens = tokens
print("-" * 30)
print(f"Total Audio Tokens Generated: {len(audio_tokens)}")
print(f"Unique Tokens: {len(set(audio_tokens))}")
print(f"First 20 Tokens: {audio_tokens[:20]}")
if len(set(audio_tokens)) <= 3 and len(audio_tokens) > 10:
print("CRITICAL: Token Collapse detected. The LLM is repeating itself.")
print("Remedy: Increase training epochs or adjust temperature/top_k.")
else:
print("Status: LLM is producing a varied token stream. The issue is likely the Decoder's projection.")
print("-" * 30)
if __name__ == "__main__":
main()