New web interface, move llm engine to separate process

tools/old_llm_engines/qwq_llm_engine.py

@@ -0,0 +1,135 @@
+# Unsloth should be imported before transformers to ensure all optimizations are applied.
+from unsloth import FastLanguageModel
+
+from pathlib import Path
+from threading import Thread
+from transformers import AutoTokenizer, TextIteratorStreamer, pipeline
+from typing import Callable, Iterator, Optional
+from xml_schema_validator import XmlLogitsProcessor
+
+from . import LlmEngine
+from .. import util
+
+class QwQLlmEngine(LlmEngine):
+    
+    def __init__(
+        self,
+        model_path: Path,
+        temperature: float,
+        xml_schema_text: Optional[str] = None,
+    ):
+        """
+        Initialize the QwQ LLM Engine.
+
+        Args:
+            model_path: Local path to the model
+            temperature: Sampling temperature
+            xml_schema_text: Optional XML schema to validate against
+        """
+        self._temperature = temperature
+
+        # Load tokenizer
+        self._tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+        )
+        
+        # Load model
+        self._model, _returned_tokenizer = FastLanguageModel.from_pretrained(
+            model_path,
+            gpu_memory_utilization = 0.5, # Reduce if out of memory
+        )
+
+        # enable unsloth optimizations
+        FastLanguageModel.for_inference(self._model)
+        
+        # Create inference pipeline
+        self._pipeline = pipeline(
+            "text-generation",
+            model=self._model,
+            tokenizer=self._tokenizer,
+            return_full_text=False,
+        )
+
+        if xml_schema_text:
+            self._logits_processor = XmlLogitsProcessor(self._tokenizer, xml_schema_text)
+        else:
+            self._logits_processor = None
+
+                        
+    def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
+        """
+        Run inference using the system prompt and main context.
+
+        Args:
+            system_prompt: The system prompt string
+            main_context: The main context string after templating
+            continuation_text: Part of the response that is already generated
+            should_stop: Callback that returns True when inference should stop
+
+        Returns:
+            Iterator[str]: An iterator that yields the generated text.
+        """
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": main_context},
+            {"role": "assistant", "content": continuation_text},
+        ]
+
+        text = self._tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False,
+        )
+
+        streamer = TextIteratorStreamer(
+            self._tokenizer,
+            skip_prompt=True,
+        )
+
+        generation_kwargs = {
+            "text_inputs": text,
+            "do_sample": True,
+            "temperature": self._temperature,
+            "streamer": streamer,
+            "use_cache": True,
+        }
+        
+        if self._logits_processor:
+            generation_kwargs["logits_processor"] = [self._logits_processor.copy()]
+
+        generation_thread = Thread(
+            target=self._pipeline,
+            kwargs=generation_kwargs
+        )
+
+        generation_thread.start()
+
+        for text in util.stop_before_value(streamer, self._tokenizer.eos_token):
+            yield text
+            if should_stop():
+                break
+
+        generation_thread.join()
+
+    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:
+        return self._pipeline.model.config.max_position_embeddings