From 1bc13e43bababc1f8eda134a87d44babffbb3866 Mon Sep 17 00:00:00 2001 From: Aleksander Majda Date: Tue, 30 Apr 2024 14:38:44 +0200 Subject: [PATCH] fix infinite loop error, add fix me loop to the chat text --- LLM_Node.py | 48 +++++++++++++++++++++++++++++++++--------------- 1 file changed, 33 insertions(+), 15 deletions(-) diff --git a/LLM_Node.py b/LLM_Node.py index b08290c..15c2eb8 100644 --- a/LLM_Node.py +++ b/LLM_Node.py @@ -72,11 +72,13 @@ class LLM_Node: def fixme(self, filename, tokenizer, model_to_use, generate_kwargs, apply_chat_template, text): self.fixmeused = True - + generated_text = "\n\n**************************************** FIX ME loop ****************************************\n\n" + generated_text += text+"\n\n\n" + file_path = os.path.join(self.comfy_ui_llm_node_path, "tmp", filename) try: - generated_text = self.generate_text(text, tokenizer, model_to_use, generate_kwargs, apply_chat_template) + generated_text += self.generate_text(text, tokenizer, model_to_use, generate_kwargs, apply_chat_template) except Exception as e: print(f"Failed to generate new text for {filename}: {e}") return @@ -93,15 +95,27 @@ class LLM_Node: with open(file_path, 'w') as file: file.write(new_code) print(f"File {filename} has been updated successfully.") + return generated_text except IOError as e: print(f"Failed to write new content to file {filename}: {e}") + return def extract_files_and_code(self, text): files_code = [] - pattern = r'\*\*(.+?\.py):\*\*\s*```python\s*([\s\S]+?)```' - matches = re.findall(pattern, text) - for filename, code in matches: - files_code.append((filename, code.strip())) + # First pattern to check for filenames and associated Python code + primary_pattern = r'\*\*(.+?\.py):\*\*\s*```python\s*([\s\S]+?)```' + matches = re.findall(primary_pattern, text) + + if matches: + for filename, code in matches: + files_code.append((filename, code.strip())) + else: + # Secondary pattern to check for Python code blocks without filenames + secondary_pattern = r'```python\s*([\s\S]+?)\s*```' + code_blocks = re.findall(secondary_pattern, text) + for code in code_blocks: + files_code.append(("main.py", code.strip())) + return files_code def write_files_to_folder(self, files_code, folder_path): @@ -139,7 +153,7 @@ class LLM_Node: generated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] else: generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True) - + print(generated_text) return generated_text def main(self, text, seed, model, max_tokens, apply_chat_template, AdvOptionsConfig=None, QuantizationConfig=None, CodingConfig=None): @@ -203,6 +217,7 @@ class LLM_Node: if config.model_type in ["t5", "gpt2", "gpt_refact", "gemma", "llama", "mistral", "qwen2"]: generated_text = self.generate_text(text, tokenizer, model_to_use, generate_kwargs, apply_chat_template) + chat_story = generated_text if CodingConfig and CodingConfig.get('execute_code'): files_code = self.extract_files_and_code(generated_text) execution_attempts = 0 @@ -223,21 +238,24 @@ class LLM_Node: env = os.environ.copy() env["SDL_VIDEODRIVER"] = "dummy" command = ['python', file_path] - subprocess.run(command, capture_output=True, text=True, check=True, env=env) # no output only check for errors - print(f"Successful Execution: {filename}") + try: + subprocess.run(command, capture_output=True, text=True, check=True, env=env, timeout=1) # no output only check for errors + print(f"Successful Execution: {filename}") + except subprocess.TimeoutExpired: + print("Execution timed out after 1 second. Terminating subprocess.") except subprocess.CalledProcessError as e: print(f"Execution failed, retrying... Error: {e.stderr}") with open(file_path, 'r') as file: file_data = file.read() text = f"Error encountered: {e.stderr}\nFix the code. Write all code. Don't take shortcut. Don't write 'same as your original code'!\n{file_data}" - self.fixme(filename, tokenizer, model_to_use, generate_kwargs, apply_chat_template, text) + chat_story += self.fixme(filename, tokenizer, model_to_use, generate_kwargs, apply_chat_template, text) execution_attempts += 1 break - else: # no more file to check + else: # no more files to check break else: # no code found break - self.log_history(text, generated_text) + self.log_history(text, chat_story) if (CodingConfig.get('project_folder')): for filename in files: source_file_path = os.path.join(tmp_folder_path, filename) @@ -247,10 +265,10 @@ class LLM_Node: subprocess.run(['python', main_py_path], capture_output=True, text=True, check=True) if os.path.exists(tmp_folder_path): shutil.rmtree(tmp_folder_path) - return (generated_text,) + return (chat_story,) else: - self.log_history(text, generated_text) - return (generated_text,) + self.log_history(text, chat_story) + return (chat_story,) elif config.model_type == "bert": return ("BERT model detected; specific task handling not implemented in this example.",)