| |
| 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) |
|
|
