File size: 3,483 Bytes
93cd35a | 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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | """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}")
# Load model architecture config from the bundled config.json
config_dir = os.path.join(os.path.dirname(__file__), "langflow")
config = LangFlowConfig.from_pretrained(config_dir)
# Build model and load weights from the downloaded safetensors file
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()
|