antonovmaxim/text-generation-webui-space

1import asyncio2import json3import sys4 5try:6    import websockets7except ImportError:8    print("Websockets package not found. Make sure it's installed.")9 10# For local streaming, the websockets are hosted without ssl - ws://11HOST = 'localhost:5005'12URI = f'ws://{HOST}/api/v1/stream'13 14# For reverse-proxied streaming, the remote will likely host with ssl - wss://15# URI = 'wss://your-uri-here.trycloudflare.com/api/v1/stream'16 17 18async def run(context):19    # Note: the selected defaults change from time to time.20    request = {21        'prompt': context,22        'max_new_tokens': 250,23        'do_sample': True,24        'temperature': 1.3,25        'top_p': 0.1,26        'typical_p': 1,27        'repetition_penalty': 1.18,28        'top_k': 40,29        'min_length': 0,30        'no_repeat_ngram_size': 0,31        'num_beams': 1,32        'penalty_alpha': 0,33        'length_penalty': 1,34        'early_stopping': False,35        'seed': -1,36        'add_bos_token': True,37        'truncation_length': 2048,38        'ban_eos_token': False,39        'skip_special_tokens': True,40        'stopping_strings': []41    }42 43    async with websockets.connect(URI, ping_interval=None) as websocket:44        await websocket.send(json.dumps(request))45 46        yield context  # Remove this if you just want to see the reply47 48        while True:49            incoming_data = await websocket.recv()50            incoming_data = json.loads(incoming_data)51 52            match incoming_data['event']:53                case 'text_stream':54                    yield incoming_data['text']55                case 'stream_end':56                    return57 58 59async def print_response_stream(prompt):60    async for response in run(prompt):61        print(response, end='')62        sys.stdout.flush()  # If we don't flush, we won't see tokens in realtime.63 64 65if __name__ == '__main__':66    prompt = "In order to make homemade bread, follow these steps:\n1)"67    asyncio.run(print_response_stream(prompt))68