File size: 1,078 Bytes
5973eb4 | 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 | """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()
|