Skip to content

Commit 79ffb78

Browse files
committed
🚧 fix(wip): mlx docker handle
1 parent 9bb44a0 commit 79ffb78

File tree

2 files changed

+12
-6
lines changed

2 files changed

+12
-6
lines changed

magnet/utils/llm/local.py

+9-2
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,13 @@
1-
import requests
1+
import platform
2+
import os
3+
os_name = platform.system()
24
from magnet.utils.globals import _f
3-
from magnet.utils.mlx import mistral
5+
if os_name == 'Darwin' and not os.getenv('DOCKER_ENV'):
6+
try:
7+
from magnet.utils.mlx import mistral
8+
except ImportError:
9+
_f("warn", "MLX module not found on macOS without DOCKER_ENV")
10+
411
import json
512

613
class LocalInference:

magnet/utils/mlx/mistral.py

+3-4
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,22 @@
11
# Copyright © 2023 Apple Inc.
22
# docstrings - 2023 Prismadic, LLC.
3-
3+
import os
44
import json
55
import time
66
from dataclasses import dataclass
77
from pathlib import Path
88
from typing import List, Optional, Tuple
99
import platform
10-
is_darwin = platform.system() == 'Darwin'
10+
os_name = platform.system()
1111

12-
if is_darwin:
12+
if os_name == 'Darwin' and not os.getenv('DOCKER_ENV'):
1313
import mlx.core as mx
1414
import mlx.nn as nn
1515
from mlx.utils import tree_unflatten
1616

1717
from sentencepiece import SentencePieceProcessor
1818
from magnet.utils.globals import _f
1919
from magnet.utils.data_classes import MistralArgs
20-
import torch.nn as nn
2120

2221
class RMSNorm(nn.Module):
2322
def __init__(self, dims: int, eps: float = 1e-5):

0 commit comments

Comments
 (0)