Fixed context usage calculation

sia/hf_llm_engine.py

@@ -1,5 +1,6 @@
-from typing import Iterator, Optional
 from huggingface_hub import InferenceClient
+from transformers import AutoTokenizer, AutoConfig
+from typing import Iterator, Optional
 
 from .llm_engine import LlmEngine
 
@@ -10,35 +11,24 @@ class HfLlmEngine(LlmEngine):
     
     def __init__(
         self,
-        model_id: str = "mistralai/Mistral-7B-Instruct-v0.2",
-        api_token: Optional[str] = None,
-        temperature: float = 0.7,
-        max_new_tokens: int = 1024,
+        model: str,
+        temperature: float,
+        api_token: Optional[str],
     ):
         """
         Initialize the HuggingFace Inference API LLM Engine.
 
         Args:
-            model_id: HuggingFace model ID to use (default: Mistral-7B-Instruct)
-            api_token: HuggingFace API token. If None, will try to read from HF_TOKEN env var
-            temperature: Sampling temperature (default: 0.7)
-            max_new_tokens: Maximum number of tokens to generate (default: 1024)
+            model: HuggingFace model ID to use
+            temperature: Sampling temperature
+            api_token: HuggingFace API token
         """
-        self.model_id = model_id
-        self.client = InferenceClient(token=api_token)
-        
-        # Generation parameters
-        self.temperature = temperature
-        self.max_new_tokens = max_new_tokens
-            
-    def set_model_path(self, model_id: str):
-        """
-        Update the model being used.
+        self._model = model
+        self._temperature = temperature
 
-        Args:
-            model_id: New HuggingFace model ID to use
-        """
-        self.model_id = model_id
+        self._tokenizer = AutoTokenizer.from_pretrained(model, token=api_token)
+        self._config = AutoConfig.from_pretrained(model, token=api_token)
+        self._client = InferenceClient(token=api_token)
                         
     def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
         """
@@ -51,21 +41,41 @@ class HfLlmEngine(LlmEngine):
         Returns:
             Iterator[str]: An iterator that yields the generated text.
         """
+        token_count=self.token_count(system_prompt, main_context)
+
         messages = [
             {"role": "system", "content": system_prompt},
             {"role": "user", "content": main_context}
         ]
         
         def stream_wrapper():
-            stream = self.client.chat_completion(
-                model=self.model_id,
+            stream = self._client.chat_completion(
+                model=self._model,
                 messages=messages,
-                temperature=self.temperature,
-                max_tokens=self.max_new_tokens,
+                temperature=self._temperature,
                 stream=True
             )
             
             for response in stream:
                 if content := response.choices[0].delta.content:
                     yield content
-        return stream_wrapper()
+        return stream_wrapper()
+
+    def token_count(self, system_prompt: str, main_context: str) -> int:
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context}
+        ]
+        prompt = self._tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        return len(self._tokenizer.encode(prompt))
+
+    def token_limit(self) -> int:
+        """
+        Get the model's context window size.
+        
+        Returns:
+            int: Maximum number of tokens the model can process
+        """
+        return self._config.max_position_embeddings