"""Call the Rust API with an ordinary text prompt (tokenizer runs in the client).

Usage:
  pip install transformers requests
  python tools/client.py --prompt 'Hello' --max-new-tokens 12

Synthetic --demo weights cannot produce meaningful text. Run with a real checkpoint.
"""
import argparse
import requests
from transformers import AutoTokenizer


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--api', default='http://127.0.0.1:8080')
    p.add_argument('--repo', default='tiiuae/Falcon-H1-0.5B-Instruct')
    p.add_argument('--prompt', default='Hello!')
    p.add_argument('--max-new-tokens', type=int, default=8)
    args = p.parse_args()
    tok = AutoTokenizer.from_pretrained(args.repo)
    ids = tok.encode(args.prompt, add_special_tokens=True)
    r = requests.post(args.api+'/v1/generate', json={'input_ids': ids,
                       'max_new_tokens': args.max_new_tokens}, timeout=1200)
    r.raise_for_status()
    answer = r.json()['output_ids']
    print(tok.decode(answer, skip_special_tokens=True))

if __name__ == '__main__':
    main()