OpenWebUI/pipelines
1
0
Fork 0
mirror of https://github.com/open-webui/pipelines synced 2025-05-31 18:48:10 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
Timothy J. Baek 2024-05-22 10:54:37 -07:00
parent a0c109e915
commit 7268268f0f
1 changed files with 14 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
pipelines/examples/mlx_pipeline.py
View File

@ -23,14 +23,18 @@ class Pipeline:
self.id = "mlx_pipeline" self.id = "mlx_pipeline"
self.name = "MLX Pipeline" self.name = "MLX Pipeline"
self.process = None self.process = None
self.model = os.getenv('MLX_MODEL', 'mistralai/Mistral-7B-Instruct-v0.2') # Default model if not set in environment variable self.model = os.getenv(
"MLX_MODEL", "mistralai/Mistral-7B-Instruct-v0.2"
) # Default model if not set in environment variable
self.port = self.find_free_port() self.port = self.find_free_port()
self.stop_sequences = os.getenv('MLX_STOP', '[INST]') # Stop sequences from environment variable self.stop_sequences = os.getenv(
"MLX_STOP", "[INST]"
) # Stop sequences from environment variable
@staticmethod @staticmethod
def find_free_port(): def find_free_port():
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(('', 0)) s.bind(("", 0))
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
return s.getsockname()[1] return s.getsockname()[1]
@ -50,9 +54,11 @@ class Pipeline:
self.process = subprocess.Popen( self.process = subprocess.Popen(
["mlx_lm.server", "--model", self.model, "--port", str(self.port)], ["mlx_lm.server", "--model", self.model, "--port", str(self.port)],
stdout=subprocess.PIPE, stdout=subprocess.PIPE,
stderr=subprocess.PIPE stderr=subprocess.PIPE,
)
print(
f"Subprocess started with PID: {self.process.pid} on port {self.port}"
) )
print(f"Subprocess started with PID: {self.process.pid} on port {self.port}")
except Exception as e: except Exception as e:
print(f"Failed to start subprocess: {e}") print(f"Failed to start subprocess: {e}")
@ -90,12 +96,12 @@ class Pipeline:
payload = { payload = {
"model": MODEL, "model": MODEL,
"messages": messages, "messages": [message.model_dump() for message in messages],
"temperature": temperature, "temperature": temperature,
"max_tokens": max_tokens, "max_tokens": max_tokens,
"top_p": top_p, "top_p": top_p,
"repetition_penalty": repetition_penalty, "repetition_penalty": repetition_penalty,
"stop": stop "stop": stop,
} }
try: try:
@ -109,4 +115,4 @@ class Pipeline:
return r.iter_lines() return r.iter_lines()
except Exception as e: except Exception as e:
return f"Error: {e}" return f"Error: {e}"