| """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() |
|
|
