File size: 1,944 Bytes
a006332 | 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 | """Command line local inference for Delta Ultra Mini without the REST API."""
from __future__ import annotations
import argparse
import logging
import os
import sys
from pathlib import Path
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
from delta.generator import DeltaGenerator
logging.basicConfig(level=os.getenv("DELTA_LOG_LEVEL", "INFO").upper())
logger = logging.getLogger(__name__)
def parse_args() -> argparse.Namespace:
"""Parse command line arguments."""
parser = argparse.ArgumentParser(description="Generate text with Delta Ultra Mini.")
parser.add_argument("--prompt", required=True, help="User prompt.")
parser.add_argument("--checkpoint_path", default="runs/delta-ultra-mini/delta_checkpoint.pt")
parser.add_argument("--tokenizer_path", default="tokenizer.json")
parser.add_argument("--config_path", default="configs/ultra_mini.json")
parser.add_argument("--max_new_tokens", type=int, default=128)
parser.add_argument("--temperature", type=float, default=0.2)
parser.add_argument("--top_k", type=int, default=20)
parser.add_argument("--top_p", type=float, default=0.9)
parser.add_argument("--repetition_penalty", type=float, default=1.08)
return parser.parse_args()
def main() -> None:
"""Load the checkpoint and print one assistant response."""
args = parse_args()
generator = DeltaGenerator.from_files(
checkpoint_path=args.checkpoint_path,
tokenizer_path=args.tokenizer_path,
config_path=args.config_path,
)
response = generator.chat(
[{"role": "user", "content": args.prompt}],
max_new_tokens=args.max_new_tokens,
temperature=args.temperature,
top_k=args.top_k,
top_p=args.top_p,
repetition_penalty=args.repetition_penalty,
)
print(response)
if __name__ == "__main__":
main()
|