| |
| """Dataset-configurable LangFlow sampler used by the DLB capture wrapper.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
| import sys |
|
|
| import torch |
| from safetensors.torch import load_file |
| from transformers import AutoTokenizer |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| UPSTREAM_ROOT = ROOT / "upstreams" / "langflow" |
| if str(UPSTREAM_ROOT) not in sys.path: |
| sys.path.insert(0, str(UPSTREAM_ROOT)) |
|
|
| from langflow import LangFlow, LangFlowConfig |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description="Generate samples with pinned LangFlow") |
| parser.add_argument("--checkpoint", required=True) |
| parser.add_argument("--num_samples", type=int, default=5) |
| parser.add_argument("--batch_size", type=int, default=1) |
| parser.add_argument("--num_steps", type=int, default=128) |
| parser.add_argument("--seq_length", type=int, default=None) |
| parser.add_argument("--seed", type=int, default=42) |
| parser.add_argument("--output") |
| parser.add_argument("--tokenizer", required=True) |
| return parser.parse_args() |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| checkpoint = Path(args.checkpoint) |
| config_dir = checkpoint.parent |
|
|
| torch.manual_seed(args.seed) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed_all(args.seed) |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| config = LangFlowConfig.from_pretrained(str(config_dir)) |
| model = LangFlow(config) |
| state_dict = load_file(str(checkpoint), device=str(device)) |
| model.load_state_dict(state_dict) |
| model = model.to(device) |
| model.eval() |
|
|
| tokenizer = AutoTokenizer.from_pretrained(args.tokenizer) |
|
|
| all_samples = [] |
| remaining = args.num_samples |
| with torch.no_grad(): |
| while remaining > 0: |
| batch = min(args.batch_size, remaining) |
| samples = model.generate_samples( |
| num_samples=batch, |
| seq_length=args.seq_length, |
| num_steps=args.num_steps, |
| device=device, |
| ) |
| all_samples.append(samples) |
| remaining -= batch |
|
|
| generated = torch.cat(all_samples, dim=0) |
| texts = tokenizer.batch_decode(generated, skip_special_tokens=True) |
| for index, text in enumerate(texts): |
| print(f"\n--- Sample {index + 1} ---") |
| print(text[:500] + ("..." if len(text) > 500 else "")) |
|
|
| if args.output is not None: |
| with open(args.output, "w", encoding="utf-8") as output_file: |
| for index, text in enumerate(texts): |
| output_file.write(f"--- Sample {index + 1} ---\n{text}\n\n") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
