mirror of
https://github.com/zenlm/zen-guard.git
synced 2026-10-30 14:58:01 +00:00
zeekay
95bf6a6295
- Zen AI branded README with proper documentation - GitHub Actions workflow for testing - Complete docs structure with paper - Training scripts and examples - Makefile for build automation
37 lines
1.2 KiB
Python
Executable File
37 lines
1.2 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Inference script for MODEL_NAME
|
|
"""
|
|
import argparse
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Run inference with MODEL_NAME")
|
|
parser.add_argument("--prompt", type=str, required=True, help="Input prompt")
|
|
parser.add_argument("--max-tokens", type=int, default=200, help="Maximum tokens to generate")
|
|
parser.add_argument("--temperature", type=float, default=0.7, help="Sampling temperature")
|
|
parser.add_argument("--model", type=str, default="zenlm/MODEL_NAME", help="Model path")
|
|
args = parser.parse_args()
|
|
|
|
print(f"Loading model: {args.model}")
|
|
model = AutoModelForCausalLM.from_pretrained(args.model)
|
|
tokenizer = AutoTokenizer.from_pretrained(args.model)
|
|
|
|
print(f"Prompt: {args.prompt}")
|
|
inputs = tokenizer(args.prompt, return_tensors="pt")
|
|
|
|
outputs = model.generate(
|
|
**inputs,
|
|
max_new_tokens=args.max_tokens,
|
|
temperature=args.temperature,
|
|
do_sample=True
|
|
)
|
|
|
|
text = tokenizer.decode(outputs[0], skip_special_tokens=True)
|
|
print(f"\nResponse:\n{text}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|