# gen_quality.py - qualitative generation samples (GPU), base LM continuation.
import os, json, torch
os.environ.setdefault('DNA_CK', '/root/dna/ckpt/base.pt')
from infer_dna import load_model, generate, load_tok

dev = 'cuda' if torch.cuda.is_available() else 'cpu'
tok = load_tok()
m, cfg = load_model(dev)
print('CONFIG', json.dumps(cfg), flush=True)

prompts = [
    "The history of science",
    "Water is important because",
    "The capital of France is",
    "Photosynthesis is the process by which",
    "Once upon a time, there was a",
    "The three primary colors are",
    "To learn how to code, you should",
    "The sun is a star that",
]
samples = []
for p in prompts:
    g = generate(m, tok, p, n=50, temp=0.0, device=dev)
    print(f'PROMPT: {p}\nCONT:   {g}\n---', flush=True)
    samples.append({'prompt': p, 'greedy': g})
json.dump(samples, open('/root/dna/quality_samples.json', 'w'), indent=2)
print('GEN_DONE', flush=True)