OpenWebUI/pipelines
1
0
Fork 0
mirror of https://github.com/open-webui/pipelines synced 2025-05-31 10:39:46 +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.name = "MLX Pipeline"
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.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
def find_free_port():
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)
return s.getsockname()[1]
@ -50,9 +54,11 @@ class Pipeline:
self.process = subprocess.Popen(
["mlx_lm.server", "--model", self.model, "--port", str(self.port)],
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:
print(f"Failed to start subprocess: {e}")
@ -90,12 +96,12 @@ class Pipeline:
payload = {
"model": MODEL,
"messages": messages,
"messages": [message.model_dump() for message in messages],
"temperature": temperature,
"max_tokens": max_tokens,
"top_p": top_p,
"repetition_penalty": repetition_penalty,
"stop": stop
"stop": stop,
}
try:
@ -109,4 +115,4 @@ class Pipeline:
return r.iter_lines()
except Exception as e:
return f"Error: {e}"
return f"Error: {e}"