Spaces:
Running on Zero
fix: Don't let a half-loaded Embedder look like a loaded one
Browse files_ensure_loaded assigned self._model from from_pretrained and only then called
.to("cuda"). When that raised on the Space, self._model was left set while
_eos_id was never reached, so the next call short-circuited on
`if self._model is not None: return` and went on to append None as a token id.
That surfaced as "RuntimeError: Could not infer dtype of NoneType" on a later
request, nowhere near the actual failure, and it is what the Space's second and
third retrieval errors were.
Everything is built into locals now and committed to self only once all of it
succeeded, so a failed load leaves the object exactly as it was and the next
attempt retries properly. A missing EOS id also raises where it happens, naming
the model and saying what depends on it, rather than travelling as a None.
Verified both backends still embed (dim 1024, unit norm) and that a load failure
leaves _model None across repeated attempts.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
- controlai_rag/embeddings.py +37 -17
|
@@ -53,33 +53,53 @@ class Embedder:
|
|
| 53 |
self._pad_id: int | None = None
|
| 54 |
|
| 55 |
def _ensure_loaded(self) -> None:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 56 |
if self._model is not None:
|
| 57 |
return
|
|
|
|
| 58 |
if self._backend == "torch":
|
| 59 |
import torch
|
| 60 |
-
from transformers import AutoModel, AutoTokenizer
|
| 61 |
-
|
| 62 |
-
self._tokenizer = AutoTokenizer.from_pretrained(self.model_id)
|
| 63 |
-
# transformers renamed torch_dtype -> dtype in 4.56; the Space
|
| 64 |
-
# resolves to whatever gradio's huggingface-hub floor allows, so
|
| 65 |
-
# detect rather than pin. See engine_torch.dtype_kwarg.
|
| 66 |
import transformers
|
|
|
|
| 67 |
|
| 68 |
-
|
| 69 |
-
|
| 70 |
-
|
| 71 |
-
|
| 72 |
-
|
|
|
|
|
|
|
|
|
|
| 73 |
)
|
| 74 |
-
|
| 75 |
-
|
| 76 |
else:
|
| 77 |
from mlx_lm import load
|
| 78 |
|
| 79 |
-
|
| 80 |
-
|
| 81 |
-
|
| 82 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 83 |
|
| 84 |
@property
|
| 85 |
def dim(self) -> int:
|
|
|
|
| 53 |
self._pad_id: int | None = None
|
| 54 |
|
| 55 |
def _ensure_loaded(self) -> None:
|
| 56 |
+
"""Load the model and tokenizer, or leave the object exactly as it was.
|
| 57 |
+
|
| 58 |
+
Everything is built into locals and committed to `self` only once all of
|
| 59 |
+
it succeeded. An earlier version assigned `self._model` from
|
| 60 |
+
`from_pretrained` and then called `.to("cuda")`, which on ZeroGPU raises:
|
| 61 |
+
`self._model` was left set, `_eos_id` was never reached, and the next
|
| 62 |
+
call short-circuited on `self._model is not None` and appended None as a
|
| 63 |
+
token id -- surfacing much later as
|
| 64 |
+
`RuntimeError: Could not infer dtype of NoneType`, nowhere near the
|
| 65 |
+
actual failure. A half-loaded embedder must not look like a loaded one.
|
| 66 |
+
"""
|
| 67 |
if self._model is not None:
|
| 68 |
return
|
| 69 |
+
|
| 70 |
if self._backend == "torch":
|
| 71 |
import torch
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 72 |
import transformers
|
| 73 |
+
from transformers import AutoModel, AutoTokenizer
|
| 74 |
|
| 75 |
+
tokenizer = AutoTokenizer.from_pretrained(self.model_id)
|
| 76 |
+
# transformers renamed torch_dtype -> dtype in 4.56, and nothing
|
| 77 |
+
# here pins a version. See engine_torch.dtype_kwarg.
|
| 78 |
+
version = tuple(int(x) for x in transformers.__version__.split(".")[:2])
|
| 79 |
+
key = "dtype" if version >= (4, 56) else "torch_dtype"
|
| 80 |
+
cuda = torch.cuda.is_available()
|
| 81 |
+
model = AutoModel.from_pretrained(
|
| 82 |
+
self.model_id, **{key: torch.float16 if cuda else torch.float32}
|
| 83 |
)
|
| 84 |
+
model = model.to("cuda" if cuda else "cpu")
|
| 85 |
+
model.eval()
|
| 86 |
else:
|
| 87 |
from mlx_lm import load
|
| 88 |
|
| 89 |
+
model, tokenizer = load(self.model_id)
|
| 90 |
+
|
| 91 |
+
ids = tokenizer.encode(EOS_TOKEN)
|
| 92 |
+
eos_id = ids[-1] if ids else tokenizer.eos_token_id
|
| 93 |
+
if eos_id is None:
|
| 94 |
+
raise RuntimeError(
|
| 95 |
+
f"{self.model_id}: could not resolve an id for {EOS_TOKEN!r}, "
|
| 96 |
+
"which last-token pooling depends on"
|
| 97 |
+
)
|
| 98 |
+
|
| 99 |
+
self._tokenizer = tokenizer
|
| 100 |
+
self._eos_id = eos_id
|
| 101 |
+
self._pad_id = tokenizer.pad_token_id or eos_id
|
| 102 |
+
self._model = model # last: this is what _ensure_loaded() checks
|
| 103 |
|
| 104 |
@property
|
| 105 |
def dim(self) -> int:
|