diff --git a/doc/server_endpoints.md b/doc/server_endpoints.md index f2679e4..2657a1b 100644 --- a/doc/server_endpoints.md +++ b/doc/server_endpoints.md @@ -288,7 +288,7 @@ Generated Events: - Parameters: - `data`: A dictionary containing the following fields: - `prompt` (string): The text prompt for text generation. - - `personality` (integer): The index of the selected personality for conditioning the text generation. + - `personality` (integer): The index of the selected personality for conditioning the text generation. If it is -1 then no personality is used and the text is assumed to be raw. - Actions: - Retrieves the selected model and client ID from the server. - Extracts the prompt and selected personality index from the request data. @@ -305,5 +305,6 @@ Generated Events: - Emits the generated text to the client through the `'text_generated'` event. Events generated: -- `'text_chunk'`: Generated text chunks are emitted to the client through this event during the text generation process. -- `'text_generated'`: Once the text generation process is complete, the final generated text is emitted to the client through this event. +- `'buzzy'`: when the server is buzzy and can't process the request, it sends this event and returns. This event have parameter `message` containing a string. +- `'text_chunk'`: Generated text chunks are emitted to the client through this event during the text generation process. The event has two parameters `chunk` and `type`. +- `'text_generated'`: Once the text generation process is complete, the final generated text is emitted to the client through this event. The event has one parameter `text` containing the full generated text. diff --git a/lollms/bindings_zoo b/lollms/bindings_zoo index c08874d..a82b707 160000 --- a/lollms/bindings_zoo +++ b/lollms/bindings_zoo @@ -1 +1 @@ -Subproject commit c08874de08eb012827aa13ab711581db9c8274b1 +Subproject commit a82b707316dbaac5b6af2bd19d531301db0ad694 diff --git a/lollms/personalities_zoo b/lollms/personalities_zoo index 14ba94d..b8c304b 160000 --- a/lollms/personalities_zoo +++ b/lollms/personalities_zoo @@ -1 +1 @@ -Subproject commit 14ba94d886bbe9713d8c4489c7ebe59317c5a057 +Subproject commit b8c304b6336d80221ad9ccd9336130ff374601c0 diff --git a/lollms/server.py b/lollms/server.py index 59c6bf6..5812e04 100644 --- a/lollms/server.py +++ b/lollms/server.py @@ -29,6 +29,8 @@ class LoLLMsServer: self.current_model = None self.personalities = [] self.answer = [''] + self.is_ready = True + @@ -88,7 +90,12 @@ class LoLLMsServer: @self.socketio.on('connect') def handle_connect(): client_id = request.sid - self.clients[client_id] = {"namespace": request.namespace, "full_discussion_blocks": []} + self.clients[client_id] = { + "namespace": request.namespace, + "full_discussion_blocks": [], + "is_generating":False, + "requested_stop":False + } ASCIIColors.success(f'Client connected with session ID: {client_id}') @self.socketio.on('disconnect') @@ -301,71 +308,116 @@ class LoLLMsServer: else: emit('personality_add_failed', {'success':False, 'error': "Personality ID not valid"}, room=request.sid) + @self.socketio.on('tokenize') + def tokenize(data): + prompt = data['prompt'] + tk = self.current_model.tokenize(prompt) + emit("tokenized", {"tokens":tk}) + + @self.socketio.on('detokenize') + def detokenize(data): + prompt = data['prompt'] + txt = self.current_model.detokenize(prompt) + emit("detokenized", {"text":txt}) + @self.socketio.on('generate_text') def handle_generate_text(data): + if not self.is_ready: + emit("buzzy", {"message":"I am buzzy. Come back later."}) + return model = self.current_model client_id = request.sid + self.clients[client_id]["is_generating"]=True + self.clients[client_id]["requested_stop"]=False prompt = data['prompt'] - personality: AIPersonality = self.personalities[data['personality']] - personality.model = model - cond_tk = personality.model.tokenize(personality.personality_conditioning) - n_cond_tk = len(cond_tk) - # Placeholder code for text generation - # Replace this with your actual text generation logic - print(f"Text generation requested by client: {client_id}") + personality_id = data['personality'] + n_predicts = data["n_predicts"] + if personality_id==-1: + # Raw text generation + print(f"Text generation requested by client: {client_id}") + self.answer[0] = '' + def callback(text, message_type: MSG_TYPE): + if message_type == MSG_TYPE.MSG_TYPE_CHUNK: + self.answer[0] = self.answer[0] + text + emit('text_chunk', {'chunk': text, 'type':MSG_TYPE.MSG_TYPE_CHUNK.value}, room=client_id) + if self.clients[client_id]["requested_stop"]: + return False + else: + return True - self.answer[0] = '' - full_discussion_blocks = self.clients[client_id]["full_discussion_blocks"] + tk = model.tokenize(prompt) + n_tokens = len(tk) + fd = model.detokenize(tk[-min(self.config.ctx_size,n_tokens):]) - if prompt != '': - if personality.processor is not None and personality.processor_cfg["process_model_input"]: - preprocessed_prompt = personality.processor.process_model_input(prompt) - else: - preprocessed_prompt = prompt + print("generating...", end="", flush=True) + generated_text = model.generate(fd, n_predict=n_predicts, callback=callback) + ASCIIColors.success(f"ok") + + # Emit the generated text to the client + emit('text_generated', {'text': generated_text}, room=client_id) + else: + personality: AIPersonality = self.personalities[personality_id] + personality.model = model + cond_tk = personality.model.tokenize(personality.personality_conditioning) + n_cond_tk = len(cond_tk) + # Placeholder code for text generation + # Replace this with your actual text generation logic + print(f"Text generation requested by client: {client_id}") + + self.answer[0] = '' + full_discussion_blocks = self.clients[client_id]["full_discussion_blocks"] + + if prompt != '': + if personality.processor is not None and personality.processor_cfg["process_model_input"]: + preprocessed_prompt = personality.processor.process_model_input(prompt) + else: + preprocessed_prompt = prompt + + if personality.processor is not None and personality.processor_cfg["custom_workflow"]: + full_discussion_blocks.append(personality.user_message_prefix) + full_discussion_blocks.append(preprocessed_prompt) + + else: + + full_discussion_blocks.append(personality.user_message_prefix) + full_discussion_blocks.append(preprocessed_prompt) + full_discussion_blocks.append(personality.link_text) + full_discussion_blocks.append(personality.ai_message_prefix) + + full_discussion = personality.personality_conditioning + ''.join(full_discussion_blocks) + + def callback(text, message_type: MSG_TYPE): + if message_type == MSG_TYPE.MSG_TYPE_CHUNK: + self.answer[0] = self.answer[0] + text + emit('text_chunk', {'chunk': text}, room=client_id) + try: + if self.clients[client_id]["requested_stop"]: + return False + else: + return True + except: # If the client is disconnected then we stop talking to it + return False + + tk = personality.model.tokenize(full_discussion) + n_tokens = len(tk) + fd = personality.model.detokenize(tk[-min(self.config.ctx_size-n_cond_tk,n_tokens):]) if personality.processor is not None and personality.processor_cfg["custom_workflow"]: - full_discussion_blocks.append(personality.user_message_prefix) - full_discussion_blocks.append(preprocessed_prompt) - + print("processing...", end="", flush=True) + generated_text = personality.processor.run_workflow(prompt, previous_discussion_text=personality.personality_conditioning+fd, callback=callback) + print(generated_text) else: + print("generating...", end="", flush=True) + generated_text = personality.model.generate(personality.personality_conditioning+fd, n_predict=personality.model_n_predicts, callback=callback) - full_discussion_blocks.append(personality.user_message_prefix) - full_discussion_blocks.append(preprocessed_prompt) - full_discussion_blocks.append(personality.link_text) - full_discussion_blocks.append(personality.ai_message_prefix) + if personality.processor is not None and personality.processor_cfg["process_model_output"]: + generated_text = personality.processor.process_model_output(generated_text) - else: - print(output.strip(),end="",flush=True) + full_discussion_blocks.append(generated_text.strip()) + print(f"{ASCIIColors.color_green}ok{ASCIIColors.color_reset}", end="", flush=True) - full_discussion = personality.personality_conditioning + ''.join(full_discussion_blocks) - - def callback(text, message_type: MSG_TYPE): - if message_type == MSG_TYPE.MSG_TYPE_CHUNK: - self.answer[0] = self.answer[0] + text - emit('text_chunk', {'chunk': text}, room=client_id) - return True - - - tk = personality.model.tokenize(full_discussion) - n_tokens = len(tk) - fd = personality.model.detokenize(tk[-min(self.config.ctx_size-n_cond_tk,n_tokens):]) - - if personality.processor is not None and personality.processor_cfg["custom_workflow"]: - print("processing...", end="", flush=True) - generated_text = personality.processor.run_workflow(prompt, previous_discussion_text=personality.personality_conditioning+fd, callback=callback) - print(generated_text) - else: - print("generating...", end="", flush=True) - generated_text = personality.model.generate(personality.personality_conditioning+fd, n_predict=personality.model_n_predicts, callback=callback) - - if personality.processor is not None and personality.processor_cfg["process_model_output"]: - generated_text = personality.processor.process_model_output(generated_text) - - full_discussion_blocks.append(generated_text.strip()) - print(f"{ASCIIColors.color_green}ok{ASCIIColors.color_reset}", end="", flush=True) - - # Emit the generated text to the client - emit('text_generated', {'text': generated_text}, room=client_id) + # Emit the generated text to the client + emit('text_generated', {'text': generated_text}, room=client_id) def build_binding(self, bindings_path: Path, cfg: LOLLMSConfig)->LLMBinding: binding_path = Path(bindings_path) / cfg["binding_name"] diff --git a/setup.py b/setup.py index 37c92dc..4526ee4 100644 --- a/setup.py +++ b/setup.py @@ -26,7 +26,7 @@ def get_all_files(path): setuptools.setup( name="lollms", - version="1.1.94", + version="1.2.0", author="Saifeddine ALOUI", author_email="aloui.saifeddine@gmail.com", description="A python library for AI personality definition", diff --git a/tests/endoints_unit_tests/example_text_gen.txt b/tests/endoints_unit_tests/example_text_gen.txt new file mode 100644 index 0000000..e4bde74 --- /dev/null +++ b/tests/endoints_unit_tests/example_text_gen.txt @@ -0,0 +1,5 @@ +Once apon a time + +In a far far galaxy, + +By the end of the summer, \ No newline at end of file diff --git a/tests/endoints_unit_tests/test_connection.py b/tests/endoints_unit_tests/test_connection.py new file mode 100644 index 0000000..ad770be --- /dev/null +++ b/tests/endoints_unit_tests/test_connection.py @@ -0,0 +1,67 @@ +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) diff --git a/tests/endoints_unit_tests/test_generation.html b/tests/endoints_unit_tests/test_generation.html new file mode 100644 index 0000000..e9baf43 --- /dev/null +++ b/tests/endoints_unit_tests/test_generation.html @@ -0,0 +1,120 @@ + + + + + + + Socket.IO Endpoint Test + + + + + +
+
+

Socket.IO Endpoint Test

+
+
+ + +
+
+ + +
+ +
+ +
+
+ + + + + +