Fix tokenizer output

requirements.txt

@@ -1,6 +1,6 @@
 transformers
-sentencepiece
 sentence-transformers
+sentencepiece
 accelerate
 faiss-cpu
 numpy

src/backend/llm/client.py

@@ -97,26 +97,31 @@ class _LocalTransformersClient(LLMClient):
         )
         self._model.eval()
 
+    def _tokenize_chat(self, messages: list[ChatMessage]):
+        """Apply the chat template and return input_ids as a (1, S) tensor."""
+        import torch
+
+        chat_messages = [{"role": m.role, "content": m.content} for m in messages]
+        token_ids = self._tokenizer.apply_chat_template(
+            chat_messages, add_generation_prompt=True,
+        )
+        return torch.tensor([token_ids], device=self._model.device)
+
     def chat(self, messages: list[ChatMessage], **kwargs) -> str:
         import torch
 
         self._ensure_loaded()
-        chat_messages = [{"role": m.role, "content": m.content} for m in messages]
-        inputs = self._tokenizer.apply_chat_template(
-            chat_messages,
-            add_generation_prompt=True,
-            return_tensors="pt",
-        ).to(self._model.device)
+        input_ids = self._tokenize_chat(messages)
         temperature = kwargs.get("temperature", self.temperature)
         with torch.no_grad():
             out = self._model.generate(
-                inputs,
+                input_ids,
                 max_new_tokens=kwargs.get("max_new_tokens", self.max_new_tokens),
                 do_sample=temperature > 0,
                 temperature=temperature if temperature > 0 else None,
                 pad_token_id=self._tokenizer.eos_token_id,
             )
-        generated = out[0][inputs.shape[-1] :]  # noqa
+        generated = out[0][input_ids.shape[-1] :]  # noqa
         return self._tokenizer.decode(generated, skip_special_tokens=True).strip()
 
     def stream(self, messages: list[ChatMessage], **kwargs):
@@ -127,16 +132,13 @@ class _LocalTransformersClient(LLMClient):
         from transformers import TextIteratorStreamer
 
         self._ensure_loaded()
-        chat_messages = [{"role": m.role, "content": m.content} for m in messages]
-        inputs = self._tokenizer.apply_chat_template(
-            chat_messages, add_generation_prompt=True, return_tensors="pt"
-        ).to(self._model.device)
+        input_ids = self._tokenize_chat(messages)
         temperature = kwargs.get("temperature", self.temperature)
         streamer = TextIteratorStreamer(
             self._tokenizer, skip_prompt=True, skip_special_tokens=True
         )
         gen_kwargs = dict(
-            input_ids=inputs,
+            input_ids=input_ids,
             streamer=streamer,
             max_new_tokens=kwargs.get("max_new_tokens", self.max_new_tokens),
             do_sample=temperature > 0,