Upload src/llm/gemma_client.py with huggingface_hub

src/llm/gemma_client.py

@@ -88,23 +88,32 @@ class GemmaClient:
             {"role": "user", "content": user_prompt},
         ]
 
-        model_inputs = self._tokenizer.apply_chat_template(
+        raw = self._tokenizer.apply_chat_template(
             messages,
             tokenize=True,
             add_generation_prompt=True,
             return_tensors="pt",
-        ).to(self._model.device)
+        )
+        # Newer transformers may return a BatchEncoding; older returns a plain tensor
+        if hasattr(raw, "input_ids"):
+            input_ids = raw.input_ids.to(self._model.device)
+            gen_kwargs = {"input_ids": input_ids}
+            if hasattr(raw, "attention_mask"):
+                gen_kwargs["attention_mask"] = raw.attention_mask.to(self._model.device)
+        else:
+            input_ids = raw.to(self._model.device)
+            gen_kwargs = {"input_ids": input_ids}
 
         with torch.no_grad():
             outputs = self._model.generate(
-                model_inputs,
+                **gen_kwargs,
                 max_new_tokens=self.config.max_new_tokens,
                 temperature=self.config.temperature,
                 do_sample=True,
                 pad_token_id=self._tokenizer.eos_token_id,
             )
 
-        prompt_len = model_inputs.shape[1]
+        prompt_len = input_ids.shape[1]
         generated_ids = outputs[0][prompt_len:]
         return self._tokenizer.decode(generated_ids, skip_special_tokens=True).strip()