| """Standalone inference script for LangFlow. |
| |
| Usage: |
| python inference.py --checkpoint /path/to/model.safetensors |
| |
| The model code lives in this repo. You only need to download the |
| safetensors checkpoint from HuggingFace — no need to clone the HF repo. |
| """ |
|
|
| import argparse |
| import os |
|
|
| import torch |
| from safetensors.torch import load_file |
| from transformers import AutoTokenizer |
|
|
| from langflow import LangFlow, LangFlowConfig |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser(description="Generate samples with LangFlow") |
| parser.add_argument( |
| "--checkpoint", type=str, required=True, |
| help="Path to the safetensors checkpoint file (e.g. model.safetensors)") |
| parser.add_argument( |
| "--num_samples", type=int, default=5, |
| help="Total number of samples to generate (default: 5)") |
| parser.add_argument( |
| "--batch_size", type=int, default=1, |
| help="Number of samples to generate per forward pass (default: 1)") |
| parser.add_argument( |
| "--num_steps", type=int, default=128, |
| help="Number of denoising steps (default: 128)") |
| parser.add_argument( |
| "--seq_length", type=int, default=1024, |
| help="Sequence length in tokens (default: 1024)") |
| parser.add_argument( |
| "--seed", type=int, default=42, |
| help="Random seed for reproducibility (default: 42)") |
| parser.add_argument( |
| "--output", type=str, default=None, |
| help="Path to save generated texts as a .txt file (optional)") |
| return parser.parse_args() |
|
|
|
|
| def main(): |
| args = parse_args() |
|
|
| 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") |
| print(f"Using device: {device}") |
|
|
| |
| config_dir = os.path.join(os.path.dirname(__file__), "langflow") |
| config = LangFlowConfig.from_pretrained(config_dir) |
|
|
| |
| print(f"Loading checkpoint: {args.checkpoint}") |
| model = LangFlow(config) |
| state_dict = load_file(args.checkpoint, device=str(device)) |
| model.load_state_dict(state_dict) |
| model = model.to(device) |
| model.eval() |
|
|
| tokenizer = AutoTokenizer.from_pretrained("gpt2") |
|
|
| print(f"Generating {args.num_samples} sample(s) " |
| f"(batch_size={args.batch_size}, steps={args.num_steps})...") |
|
|
| 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 |
|
|
| all_samples = torch.cat(all_samples, dim=0) |
| texts = tokenizer.batch_decode(all_samples, skip_special_tokens=True) |
| for i, text in enumerate(texts): |
| print(f"\n--- Sample {i + 1} ---") |
| print(text[:500] + ("..." if len(text) > 500 else "")) |
|
|
| if args.output is not None: |
| with open(args.output, "w", encoding="utf-8") as f: |
| for i, text in enumerate(texts): |
| f.write(f"--- Sample {i + 1} ---\n{text}\n\n") |
| print(f"\nSaved {len(texts)} sample(s) to {args.output}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
