""" oss_assistant.py ~~~~~~~~~~~~~~~~ Production-ready OSS AI assistant backed by Hugging Face Transformers. Model : Qwen/Qwen2.5-0.5B-Instruct (CPU-compatible, ~1 GB download) Features: - Multi-turn conversation via ConversationMemory (sliding-window + token budget) - Token-overflow guard: oldest turns evicted before every inference call - Configurable temperature, top-p, repetition penalty, max new tokens - Lazy model loading (model loaded once, reused across calls) - Thread-safe inference via a module-level lock - Streamlit-compatible (returns plain strings, no blocking I/O) - Full error handling and loguru logging Usage (standalone): from models.oss_assistant import OSSAssistant bot = OSSAssistant() print(bot.chat("Hello! What can you do?")) Usage (Streamlit): if "oss_bot" not in st.session_state: st.session_state.oss_bot = OSSAssistant() response = st.session_state.oss_bot.chat(user_input) """ from __future__ import annotations import threading import time from dataclasses import dataclass, field from typing import List, Dict, Optional import torch from transformers import ( AutoTokenizer, AutoModelForCausalLM, GenerationConfig, TextIteratorStreamer, ) from models.base_assistant import BaseAssistant from models.memory_manager import ConversationMemory from dotenv import load_dotenv from models.logger_config import logger load_dotenv() # ── Module-level lock — prevents concurrent inference on a shared model ─────── _INFERENCE_LOCK = threading.Lock() # ───────────────────────────────────────────────────────────────────────────── # Configuration dataclass # ───────────────────────────────────────────────────────────────────────────── @dataclass class AssistantConfig: """ All tunable parameters for the OSS assistant. Changing any field on an already-instantiated OSSAssistant takes effect on the very next call to chat(). """ # Model identity model_id: str = "Qwen/Qwen2.5-0.5B-Instruct" # Generation knobs max_new_tokens: int = 512 # tokens the model may produce per turn temperature: float = 0.7 # higher → more creative; lower → more factual top_p: float = 0.9 # nucleus sampling threshold repetition_penalty: float = 1.1 # penalise repeated tokens # Context / memory management max_history_turns: int = 10 # sliding window: keep last N human/AI pairs max_context_tokens: int = 1_800 # token budget for system + history before inference system_prompt: str = ( "You are a helpful, harmless, and honest AI assistant. " "Answer concisely and accurately." ) # Runtime device: Optional[str] = None # None → auto-detect (cuda > mps > cpu) torch_dtype: Optional[str] = None # None → float32 on CPU, float16 on GPU # ───────────────────────────────────────────────────────────────────────────── # Main assistant class # ───────────────────────────────────────────────────────────────────────────── class OSSAssistant(BaseAssistant): """ Multi-turn conversational assistant using a locally-run HuggingFace model. Implements the BaseAssistant interface — interchangeable with GroqAssistant. The class is responsible for: 1. Loading the tokenizer and model (once, lazily). 2. Delegating memory/history management to BaseAssistant / ConversationMemory. 3. Running inference with configurable sampling parameters. 4. Returning the plain assistant reply as a string. """ def __init__(self, config: Optional[AssistantConfig] = None) -> None: self.config = config or AssistantConfig() # ConversationMemory owns all history, windowing, and token-budget logic. self._memory = ConversationMemory( max_turns=self.config.max_history_turns, max_context_tokens=self.config.max_context_tokens, system_prompt=self.config.system_prompt, ) # Lazy-loaded — set by _ensure_model_loaded() self._tokenizer: Optional[AutoTokenizer] = None self._model: Optional[AutoModelForCausalLM] = None self._device: Optional[str] = None logger.info(f"OSSAssistant initialised | model={self.config.model_id}") # ── Abstract interface implementation ───────────────────────────────────── @property def backend_name(self) -> str: """Model identifier used in get_info() and UI labels.""" return self.config.model_id def chat(self, user_message: str) -> str: """ Send a user message and return the assistant's reply. Flow ---- 1. Validate input. 2. Load model (once, lazily). 3. Register user message in memory (opens a new Turn). 4. _run_inference() enforces token budget before tokenising. 5. On success: save assistant reply; on failure: rollback the open Turn. """ if not user_message or not user_message.strip(): return "Please enter a message." self._sync_memory_config() self._ensure_model_loaded() self._memory.add_user_message(user_message) try: reply = self._run_inference() except Exception as exc: logger.error(f"Inference failed: {exc}") self._memory.rollback_last_user_message() return f"[Error during inference: {exc}]" self._memory.add_assistant_reply(reply) logger.debug(f"Turn complete | user={user_message[:60]!r} | reply={reply[:60]!r}") return reply def stream(self, user_message: str) -> Iterator[str]: """ Yield the assistant's reply token-by-token using TextIteratorStreamer. A background daemon thread runs model.generate() and feeds tokens into a TextIteratorStreamer queue. The main thread yields from the queue, making this compatible with Streamlit's st.write_stream(). Memory is updated after the full reply is assembled (atomic save). On error the open turn is rolled back. """ import threading user_message = (user_message or "").strip() if not user_message: yield "Please enter a message." return self._sync_memory_config() self._ensure_model_loaded() self._memory.add_user_message(user_message) try: with _INFERENCE_LOCK: # Enforce token budget before building the prompt self._memory.enforce_token_budget(self._tokenizer) messages = self._memory.as_prompt_messages() prompt_text = self._tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) inputs = self._tokenizer( prompt_text, return_tensors="pt", truncation=True, max_length=self.config.max_context_tokens, ).to(self._model.device) input_len = inputs["input_ids"].shape[1] # TextIteratorStreamer lets us yield tokens from the main thread # while generate() runs in a background thread. streamer = TextIteratorStreamer( self._tokenizer, skip_prompt=True, # don't echo the prompt back skip_special_tokens=True, ) gen_config = GenerationConfig( max_new_tokens=self.config.max_new_tokens, do_sample=self.config.temperature > 0, temperature=self.config.temperature if self.config.temperature > 0 else None, top_p=self.config.top_p, repetition_penalty=self.config.repetition_penalty, pad_token_id=self._tokenizer.pad_token_id, eos_token_id=self._tokenizer.eos_token_id, ) # Launch generation in a daemon thread so it doesn't block gen_thread = threading.Thread( target=self._model.generate, kwargs={**inputs, "generation_config": gen_config, "streamer": streamer}, daemon=True, ) gen_thread.start() # Collect and yield tokens as they arrive reply_parts: list[str] = [] for token_text in streamer: reply_parts.append(token_text) yield token_text gen_thread.join() except Exception as exc: logger.error(f"Stream inference failed: {exc}") self._memory.rollback_last_user_message() yield f"[Error during inference: {exc}]" return full_reply = "".join(reply_parts).strip() if full_reply: self._memory.add_assistant_reply(full_reply) else: self._memory.rollback_last_user_message() logger.warning("OSS model produced an empty stream reply; turn rolled back.") @property def is_loaded(self) -> bool: """True once the model has been loaded into memory.""" return self._model is not None # ── Model loading ───────────────────────────────────────────────────────── def _ensure_model_loaded(self) -> None: """Load tokenizer + model exactly once. Thread-safe via the module lock.""" if self._model is not None: return # already loaded with _INFERENCE_LOCK: # Double-check inside the lock (another thread may have loaded first) if self._model is not None: return logger.info(f"Loading model: {self.config.model_id}") t0 = time.perf_counter() # ── Resolve device ───────────────────────────────────────────── if self.config.device: self._device = self.config.device elif torch.cuda.is_available(): self._device = "cuda" elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): self._device = "mps" else: self._device = "cpu" logger.info(f"Inference device: {self._device.upper()}") # ── Resolve dtype ────────────────────────────────────────────── if self.config.torch_dtype: dtype = getattr(torch, self.config.torch_dtype) elif self._device == "cpu": dtype = torch.float32 # float16 is slower on CPU else: dtype = torch.float16 # saves VRAM on GPU # ── Load tokenizer ───────────────────────────────────────────── self._tokenizer = AutoTokenizer.from_pretrained( self.config.model_id, trust_remote_code=True, # required for Qwen models ) # Ensure pad token is set (Qwen may not set it by default) if self._tokenizer.pad_token is None: self._tokenizer.pad_token = self._tokenizer.eos_token # ── Load model ───────────────────────────────────────────────── self._model = AutoModelForCausalLM.from_pretrained( self.config.model_id, torch_dtype=dtype, device_map="auto", # auto-places layers on GPU/CPU low_cpu_mem_usage=True, # stream weights instead of double-buffering trust_remote_code=True, ) self._model.eval() # disable dropout for deterministic inference elapsed = time.perf_counter() - t0 logger.success(f"Model loaded in {elapsed:.1f}s on {self._device.upper()}") # ── Config sync ─────────────────────────────────────────────────────────── def _sync_memory_config(self) -> None: """ Push live AssistantConfig values into the ConversationMemory object. Called at the top of chat() so Streamlit sidebar changes take effect immediately without requiring a session restart. """ self._memory.max_turns = self.config.max_history_turns self._memory.max_context_tokens = self.config.max_context_tokens self._memory.system_prompt = self.config.system_prompt # ── Inference ───────────────────────────────────────────────────────────── def _run_inference(self) -> str: """ Enforce token budget, tokenize, run generation, decode reply. Steps ----- 1. enforce_token_budget() — evicts oldest complete turns until the rendered prompt fits within max_context_tokens. 2. apply_chat_template() — renders history to a single string using the model's native template (handles all special tokens). 3. Tokenize with explicit truncation as a last-resort safety net. 4. Generate with per-call GenerationConfig (picks up sidebar changes). 5. Slice prompt tokens off the output; decode only new tokens. Uses the module-level lock — only one thread runs inference at a time. """ with _INFERENCE_LOCK: # ── Step 1: token-budget enforcement ───────────────────────────── # Evicts oldest complete turns so the prompt fits in context. final_token_count = self._memory.enforce_token_budget(self._tokenizer) logger.debug(f"Prompt tokens after budget enforcement: {final_token_count}") # ── Step 2: render prompt via chat template ─────────────────────── messages = self._memory.as_prompt_messages() prompt_text = self._tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, # appends the "<|im_start|>assistant" trigger ) # ── Step 3: tokenize (truncation = last-resort safety net) ──────── inputs = self._tokenizer( prompt_text, return_tensors="pt", padding=True, truncation=True, max_length=self.config.max_context_tokens, ).to(self._model.device) input_len = inputs["input_ids"].shape[1] logger.debug(f"Tokenized prompt length: {input_len} tokens") # ── Step 4: build GenerationConfig per-call ─────────────────────── # Rebuilding each call means Streamlit slider changes take effect # immediately without restarting the session. gen_config = GenerationConfig( max_new_tokens=self.config.max_new_tokens, do_sample=self.config.temperature > 0, temperature=self.config.temperature if self.config.temperature > 0 else None, top_p=self.config.top_p, repetition_penalty=self.config.repetition_penalty, pad_token_id=self._tokenizer.pad_token_id, eos_token_id=self._tokenizer.eos_token_id, ) # ── Step 5: generate + decode ───────────────────────────────────── with torch.inference_mode(): # disables grad tracking — faster + less RAM output_ids = self._model.generate( **inputs, generation_config=gen_config, ) # Keep only the newly generated token ids (strip the prompt prefix) new_token_ids = output_ids[0][input_len:] reply = self._tokenizer.decode(new_token_ids, skip_special_tokens=True) logger.debug(f"Generated {len(new_token_ids)} new token(s)") return reply.strip() # ── Dunder helpers ──────────────────────────────────────────────────────── def __repr__(self) -> str: status = "loaded" if self.is_loaded else "not loaded" return ( f"OSSAssistant(model={self.config.model_id!r}, " f"device={self._device!r}, status={status}, turns={self.turn_count})" )