main.py 792 B

1234567891011121314151617181920212223242526272829303132
  1. import http.client
  2. import json
  3. import os
  4. import sys
  5. if len(sys.argv) < 2:
  6. print("Usage: python main.py <model file>")
  7. sys.exit(1)
  8. conn = http.client.HTTPConnection('localhost', 11434)
  9. headers = { 'Content-Type': 'application/json' }
  10. # generate text from the model
  11. conn.request("POST", "/api/generate", json.dumps({
  12. 'model': os.path.join(os.getcwd(), sys.argv[1]),
  13. 'prompt': 'write me a short story',
  14. 'stream': True
  15. }), headers)
  16. response = conn.getresponse()
  17. def parse_generate(data):
  18. for event in data.decode('utf-8').split("\n"):
  19. if not event:
  20. continue
  21. yield event
  22. if response.status == 200:
  23. for chunk in response:
  24. for event in parse_generate(chunk):
  25. print(json.loads(event)['response'], end="", flush=True)