api-example.py 1.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960
  1. '''
  2. This is an example on how to use the API for oobabooga/text-generation-webui.
  3. Make sure to start the web UI with the following flags:
  4. python server.py --model MODEL --listen --no-stream
  5. Optionally, you can also add the --share flag to generate a public gradio URL,
  6. allowing you to use the API remotely.
  7. '''
  8. import requests
  9. # Server address
  10. server = "127.0.0.1"
  11. # Generation parameters
  12. # Reference: https://huggingface.co/docs/transformers/main_classes/text_generation#transformers.GenerationConfig
  13. params = {
  14. 'max_new_tokens': 200,
  15. 'do_sample': True,
  16. 'temperature': 0.5,
  17. 'top_p': 0.9,
  18. 'typical_p': 1,
  19. 'repetition_penalty': 1.05,
  20. 'top_k': 0,
  21. 'min_length': 0,
  22. 'no_repeat_ngram_size': 0,
  23. 'num_beams': 1,
  24. 'penalty_alpha': 0,
  25. 'length_penalty': 1,
  26. 'early_stopping': False,
  27. }
  28. # Input prompt
  29. prompt = "What I would like to say is the following: "
  30. response = requests.post(f"http://{server}:7860/run/textgen", json={
  31. "data": [
  32. prompt,
  33. params['max_new_tokens'],
  34. params['do_sample'],
  35. params['max_new_tokens'],
  36. params['temperature'],
  37. params['top_p'],
  38. params['typical_p'],
  39. params['repetition_penalty'],
  40. params['top_k'],
  41. params['min_length'],
  42. params['no_repeat_ngram_size'],
  43. params['num_beams'],
  44. params['penalty_alpha'],
  45. params['length_penalty'],
  46. params['early_stopping'],
  47. ]
  48. }).json()
  49. reply = response["data"][0]
  50. print(reply)