mirror of
https://github.com/ParisNeo/lollms.git
synced 2024-12-18 20:27:58 +00:00
68 lines
2.2 KiB
Python
68 lines
2.2 KiB
Python
|
import argparse
|
||
|
|
import socketio
|
||
|
|
from pathlib import Path
|
||
|
|
from lollms import MSG_TYPE
|
||
|
|
import time
|
||
|
|
|
||
|
|
# Connect to the Socket.IO server
|
||
|
|
sio = socketio.Client()
|
||
|
|
|
||
|
|
|
||
|
|
# Event handler for receiving generated text
|
||
|
|
@sio.event
|
||
|
|
def text_generated(data):
|
||
|
|
print('Generated text:', data)
|
||
|
|
|
||
|
|
def test_generate_text(host, port, text_file):
|
||
|
|
# Read the text file and split by multiple newlines
|
||
|
|
print("Loading file")
|
||
|
|
with open(text_file, 'r') as file:
|
||
|
|
prompts = file.read().split('\n\n')
|
||
|
|
|
||
|
|
is_ready=[False]
|
||
|
|
# Event handler for successful connection
|
||
|
|
@sio.event
|
||
|
|
def connect():
|
||
|
|
print('Connected to Socket.IO server')
|
||
|
|
for prompt in prompts:
|
||
|
|
if prompt:
|
||
|
|
# Trigger the 'generate_text' event with the prompt
|
||
|
|
is_ready[0]=False
|
||
|
|
print(f"Sending prompt:{prompt}")
|
||
|
|
sio.emit('generate_text', {'prompt': prompt, 'personality':-1, "n_predicts":1024})
|
||
|
|
while is_ready[0]==False:
|
||
|
|
time.sleep(0.1)
|
||
|
|
|
||
|
|
@sio.event
|
||
|
|
def text_chunk(data):
|
||
|
|
print(data["chunk"],end="",flush=True)
|
||
|
|
|
||
|
|
@sio.event
|
||
|
|
def text_generated(data):
|
||
|
|
print("text_generated_ok")
|
||
|
|
print(data["text"])
|
||
|
|
is_ready[0]=True
|
||
|
|
|
||
|
|
print(f"Connecting to http://{host}:{port}")
|
||
|
|
# Connect to the Socket.IO server
|
||
|
|
sio.connect(f'http://{host}:{port}')
|
||
|
|
|
||
|
|
# Start the event loop
|
||
|
|
sio.wait()
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
# Parse command-line arguments
|
||
|
|
parser = argparse.ArgumentParser(description='Socket.IO endpoint test')
|
||
|
|
parser.add_argument('--host', type=str, default='localhost', help='Socket.IO server host')
|
||
|
|
parser.add_argument('--port', type=int, default=9600, help='Socket.IO server port')
|
||
|
|
parser.add_argument('--text-file', type=str, default=str(Path(__file__).parent/"example_text_gen.txt"),help='Path to the text file')
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
# Verify if the text file exists
|
||
|
|
text_file_path = Path(args.text_file)
|
||
|
|
if not text_file_path.is_file():
|
||
|
|
print(f"Error: The provided text file '{args.text_file}' does not exist.")
|
||
|
|
else:
|
||
|
|
# Run the test with provided arguments
|
||
|
|
test_generate_text(args.host, args.port, args.text_file)
|