Fixed context usage calculation

sia/local_llm_engine.py

@@ -1,5 +1,5 @@
 from threading import Thread
-from typing import Iterator
+from typing import Iterator, Optional
 
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
 import torch
@@ -8,23 +8,23 @@ from . import util
 from .llm_engine import LlmEngine
 
 class LocalLlmEngine(LlmEngine):
-    def __init__(self, model_path: str):
+    def __init__(
+        self,
+        model_path: str,
+        temperature: float,
+        token_limit: int,
+    ):
         """
         Initialize the LLM Engine with a model path.
 
         Args:
             model_path: Path to the model weights to be used.
+            temperature: Temperature for sampling
+            token_limit: Maximum number of tokens to generate
         """
-        self.set_model_path(model_path)
-
-    def set_model_path(self, model_path: str):
-        """
-        Load the model from the specified path.
-
-        Args:
-            model_path: Path to the model weights to load.
-        """
-        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self._temperature = temperature
+        self._token_limit = token_limit
+        self._tokenizer = AutoTokenizer.from_pretrained(model_path)
         model = AutoModelForCausalLM.from_pretrained(
             model_path,
             return_dict=True,
@@ -33,14 +33,14 @@ class LocalLlmEngine(LlmEngine):
             device_map="auto",
             trust_remote_code=True,
         )
-        if self.tokenizer.pad_token_id is None:
-            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+        if self._tokenizer.pad_token_id is None:
+            self._tokenizer.pad_token_id = self._tokenizer.eos_token_id
         if model.config.pad_token_id is None:
             model.config.pad_token_id = model.config.eos_token_id
-        self.pipeline = pipeline(
+        self._pipeline = pipeline(
             "text-generation",
             model=model,
-            tokenizer=self.tokenizer,
+            tokenizer=self._tokenizer,
             torch_dtype=torch.bfloat16,
             device_map="auto",
             return_full_text=False,
@@ -61,19 +61,49 @@ class LocalLlmEngine(LlmEngine):
             {"role": "system", "content": system_prompt},
             {"role": "user", "content": main_context}
         ]
-        prompt = self.tokenizer.apply_chat_template(
+        prompt = self._tokenizer.apply_chat_template(
             messages, tokenize=False, add_generation_prompt=True
         )
         streamer = TextIteratorStreamer(
-            self.tokenizer,
+            self._tokenizer,
             skip_prompt=True
         )
         pipeline_kwargs = dict(
             text_inputs=prompt,
             do_sample=True,
-            max_new_tokens=1024,
+            temperature=self._temperature,
+            max_new_tokens=self._token_limit,
             streamer=streamer
         )
-        thread = Thread(target=self.pipeline, kwargs=pipeline_kwargs)
+        thread = Thread(target=self._pipeline, kwargs=pipeline_kwargs)
         thread.start()
         return util.stop_before_value(streamer, '<|eot_id|>')
+
+    def token_count(self, system_prompt: str, main_context: str) -> int:
+        """
+        Count tokens for the given system prompt and main context.
+        
+        Args:
+            system_prompt: The system prompt string
+            main_context: The main context string
+            
+        Returns:
+            int: Total number of tokens
+        """
+        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._pipeline.model.config.max_position_embeddings