File size: 2,735 Bytes
23ea6bd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | """
Script 3: Interactive Text Generation & Inference for Nova 1.0
"""
import os
import argparse
import torch
from config.model_config import Nova1Config
from src.tokenizer.bpe_tokenizer import Nova1Tokenizer
from src.model.nova1_hrm import Nova1HRM
from src.inference.generate import Nova1Generator
def main():
parser = argparse.ArgumentParser(description="Generate text with Nova 1.0")
parser.add_argument("--checkpoint", type=str, default="checkpoints/nova1_final.pt", help="Path to Nova 1.0 checkpoint")
default_tok = "checkpoints/nova1_gemini_tokenizer.json" if os.path.exists("checkpoints/nova1_gemini_tokenizer.json") else "checkpoints/nova1_tokenizer.json"
parser.add_argument("--tokenizer_path", type=str, default=default_tok, help="Path to tokenizer JSON")
parser.add_argument("--prompt", type=str, default="Once upon a time in a small village,", help="Generation prompt")
parser.add_argument("--max_tokens", type=int, default=60, help="Max new tokens to generate")
parser.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature")
parser.add_argument("--top_k", type=int, default=40, help="Top-K sampling")
parser.add_argument("--top_p", type=float, default=0.9, help="Top-P nucleus sampling")
parser.add_argument("--N_cycles", type=int, default=None, help="Test-time reasoning depth N cycles")
parser.add_argument("--T_steps", type=int, default=None, help="Test-time low-level steps T")
args = parser.parse_args()
if not os.path.exists(args.tokenizer_path):
print(f"Error: Tokenizer not found at '{args.tokenizer_path}'.")
return
tokenizer = Nova1Tokenizer.load(args.tokenizer_path)
if not os.path.exists(args.checkpoint):
print(f"Error: Model checkpoint not found at '{args.checkpoint}'.")
return
checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
config: Nova1Config = checkpoint.get("config", Nova1Config(vocab_size=tokenizer.vocab_size))
model = Nova1HRM(config)
model.load_state_dict(checkpoint["model_state"])
print(f"Loaded Nova 1.0 model checkpoint from '{args.checkpoint}'.")
generator = Nova1Generator(model=model, tokenizer=tokenizer, config=config)
print(f"\n--- Nova 1.0 Text Generation ---")
print(f"Prompt: {args.prompt}")
output_text = generator.generate(
prompt=args.prompt,
max_new_tokens=args.max_tokens,
temperature=args.temperature,
top_k=args.top_k,
top_p=args.top_p,
N_cycles=args.N_cycles,
T_steps=args.T_steps
)
print("\n--- Output ---")
print(output_text)
print("--------------\n")
if __name__ == "__main__":
main()
|