From 83aa6107a17e62e01a64e8a79cac849a7f155595 Mon Sep 17 00:00:00 2001 From: Roman Kuzmenko Date: Thu, 17 Oct 2024 21:42:30 -0700 Subject: [PATCH] Added suppotr for Llama-3.2, and images for all olllama models --- partcad/src/partcad/ai.py | 2 +- partcad/src/partcad/ai_ollama.py | 20 +++++++++- .../src/partcad/part_factory_feature_ai.py | 38 +++++++++++++------ 3 files changed, 45 insertions(+), 15 deletions(-) diff --git a/partcad/src/partcad/ai.py b/partcad/src/partcad/ai.py index 95b40ed6..3976e176 100644 --- a/partcad/src/partcad/ai.py +++ b/partcad/src/partcad/ai.py @@ -29,7 +29,7 @@ "gemini-pro-vision", "gemini-1.5-pro", "gemini-1.5-flash", - "llama3.1*", + "llama3.*", "codellama*", "codegemma*", "gemma*", diff --git a/partcad/src/partcad/ai_ollama.py b/partcad/src/partcad/ai_ollama.py index 5f0082b1..27be05bb 100644 --- a/partcad/src/partcad/ai_ollama.py +++ b/partcad/src/partcad/ai_ollama.py @@ -7,8 +7,11 @@ # Licensed under Apache License, Version 2.0. # +import base64 import importlib import httpx +from path import Path +import re import threading import time from typing import Any @@ -70,8 +73,18 @@ def generate_ollama( if not ollama_once(): return None - if "INSERT_IMAGE_HERE" in prompt: - raise NotImplementedError("Images are not supported by Ollama") + image_content = [] + + def insert_image(match): + filename = match.group(1) + image_index = len(image_content) + image_content.append( + Path(filename).read_bytes(), + # base64.b64encode(Path(filename).read_bytes()).decode(), + ) + return f"The attached image number {image_index}.\n" + + prompt = re.sub(r"INSERT_IMAGE_HERE\(([^)]*)\)", insert_image, prompt) if "tokens" in config: tokens = config["tokens"] @@ -106,10 +119,13 @@ def generate_ollama( top_k=top_k, temperature=temperature, ) + pc_logging.debug("Prompt: %s" % prompt) + pc_logging.debug("Images: %d" % len(image_content)) response = ollama.generate( model=model, context=[], # do not accumulate context uncontrollably prompt=prompt, + images=image_content, options=options, ) except httpx.ConnectError as e: diff --git a/partcad/src/partcad/part_factory_feature_ai.py b/partcad/src/partcad/part_factory_feature_ai.py index c8fc0520..793753a4 100644 --- a/partcad/src/partcad/part_factory_feature_ai.py +++ b/partcad/src/partcad/part_factory_feature_ai.py @@ -102,7 +102,9 @@ def on_init_ai(self): constructor to finalize the AI initialization. At the time of the call self.part and self.instantiate must be already defined.""" self.part.generate = lambda path: self._create_file(path) - self.part.change = lambda path, change=None: self._change_file(path, change) + self.part.change = lambda path, change=None: self._change_file( + path, change + ) # If uncommented out, this makes the package initialization # unaccceptably slow @@ -199,10 +201,8 @@ def _create_file(self, path): # Validate the image by rendering it, # attempt to correct the script if rendering doesn't work - image_filename, changed_script = ( - self._validate_and_fix( - changed_script, candidate_id - ) + image_filename, changed_script = self._validate_and_fix( + changed_script, candidate_id ) # Check if the model was valid if image_filename is not None: @@ -239,6 +239,11 @@ def _change_file(self, path, change=None): script = open(path, "r").read() + if change is not None: + change = change.strip() + if change == "": + change = None + image_filename, error_text = self._render_image(script, 0) if image_filename is None or error_text: pc_logging.error( @@ -251,7 +256,9 @@ def _change_file(self, path, change=None): # Attempt to change the script once more by comparing the result with # the original request - changed_scripts = self._change_script(None, script, image_filename, change) + changed_scripts = self._change_script( + None, script, image_filename, change + ) for changed_script in changed_scripts: pc_logging.debug( "Generated the changed script: %s" % changed_script @@ -278,7 +285,9 @@ def _change_file(self, path, change=None): new_script = script_candidates[0][1] else: # Compare the images and select the best one - new_script = self.select_best_image(script_candidates, change=change) + new_script = self.select_best_image( + script_candidates, change=change + ) if new_script == script: pc_logging.info("The script was not changed") @@ -365,7 +374,7 @@ def _generate_script(self, csg_instructions): """This method generates a script given specific CSG description.""" prompt = """You are an AI assistant in an engineering department. -You are helping engineers to create programmatic scripts that produce CAD geometry data +You are helping engineers by writing scripts that produce CAD geometry data for parts, mechanisms, buildings or anything else. The scripts you create are fully functional and can be used right away, as is, in automated workflows. Assume that the scripts you produce are used automatically to render 3D models and to validate them. @@ -415,7 +424,9 @@ def _generate_script(self, csg_instructions): return scripts - def _change_script(self, csg_instructions, script, rendered_image, change=None): + def _change_script( + self, csg_instructions, script, rendered_image, change=None + ): """This method changes the script given the original request and the produced script.""" config = copy.copy(self.ai_config) @@ -469,7 +480,7 @@ def _change_script(self, csg_instructions, script, rendered_image, change=None): if change is not None: prompt += f"\n\n{change}\n" else: - prompt += """ + prompt += """ Please, analyze whether the produced script and image match the original request (where the original image and description take precedence @@ -753,12 +764,15 @@ def select_best_image(self, script_candidates, change=None): prompt += "INSERT_IMAGE_HERE(%s)\n" % image_filename if change is not None: - prompt += """ + prompt += ( + """ Subsequently, the following changes were requested (until "CHANGE END"): %s CHANGE END -""" % change +""" + % change + ) prompt += """