atakan Claude Opus 5 commited on
Commit
929b04a
·
1 Parent(s): ac301c2

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>

Files changed (1) hide show
  1. controlai_rag/embeddings.py +37 -17
controlai_rag/embeddings.py CHANGED
@@ -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
- _v = tuple(int(x) for x in transformers.__version__.split(".")[:2])
69
- _key = "dtype" if _v >= (4, 56) else "torch_dtype"
70
- self._model = AutoModel.from_pretrained(
71
- self.model_id,
72
- **{_key: torch.float16 if torch.cuda.is_available() else torch.float32},
 
 
 
73
  )
74
- self._model = self._model.to("cuda" if torch.cuda.is_available() else "cpu")
75
- self._model.eval()
76
  else:
77
  from mlx_lm import load
78
 
79
- self._model, self._tokenizer = load(self.model_id)
80
- ids = self._tokenizer.encode(EOS_TOKEN)
81
- self._eos_id = ids[-1] if ids else self._tokenizer.eos_token_id
82
- self._pad_id = self._tokenizer.pad_token_id or self._eos_id
 
 
 
 
 
 
 
 
 
 
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: