inference working

packages/node/meshnet_node/model_backend.py

@@ -241,6 +241,24 @@ class TorchModelShard:
             yield token_text
         t.join()
 
+    def count_prompt_tokens(self, messages: list[dict]) -> int:
+        """Return tokenizer-backed prompt token count for OpenAI usage metadata."""
+        encoded = self._encode_messages(messages)
+        input_ids = encoded["input_ids"]
+        return int(input_ids.shape[-1])
+
+    def count_text_tokens(self, text: str) -> int:
+        """Return tokenizer-backed completion token count for OpenAI usage metadata."""
+        try:
+            encoded = self.tokenizer(
+                text,
+                return_tensors="pt",
+                add_special_tokens=False,
+            )
+        except TypeError:
+            encoded = self.tokenizer(text, return_tensors="pt")
+        return int(encoded["input_ids"].shape[-1])
+
     def _encode_messages(self, messages: list[dict]) -> dict:
         """Format messages with chat template (if available) and tokenize."""
         if hasattr(self.tokenizer, "apply_chat_template"):