Style improvements (#1957)

This commit is contained in:
oobabooga 2023-05-09 22:49:39 -03:00 committed by GitHub
parent 334486f527
commit 3913155c1f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
23 changed files with 64 additions and 50 deletions

View file

@ -5,7 +5,7 @@ import sys
try:
import websockets
except ImportError:
print("Websockets package not found. Make sure it's installed.")
print("Websockets package not found. Make sure it's installed.")
# For local streaming, the websockets are hosted without ssl - ws://
HOST = 'localhost:5005'
@ -14,6 +14,7 @@ URI = f'ws://{HOST}/api/v1/stream'
# For reverse-proxied streaming, the remote will likely host with ssl - wss://
# URI = 'wss://your-uri-here.trycloudflare.com/api/v1/stream'
async def run(context):
# Note: the selected defaults change from time to time.
request = {
@ -42,7 +43,7 @@ async def run(context):
async with websockets.connect(URI, ping_interval=None) as websocket:
await websocket.send(json.dumps(request))
yield context # Remove this if you just want to see the reply
yield context # Remove this if you just want to see the reply
while True:
incoming_data = await websocket.recv()
@ -58,7 +59,7 @@ async def run(context):
async def print_response_stream(prompt):
async for response in run(prompt):
print(response, end='')
sys.stdout.flush() # If we don't flush, we won't see tokens in realtime.
sys.stdout.flush() # If we don't flush, we won't see tokens in realtime.
if __name__ == '__main__':