Spaces:
Sleeping
Sleeping
| from typing import Dict, Any, Optional | |
| from src.intent import detect_intent | |
| from src.templates import TEMPLATES | |
| DEFAULT_GEN_ARGS = { | |
| "max_tokens": 300, | |
| "temperature": 0.7, | |
| "top_p": 0.95 | |
| } | |
| MSG_SEPARATOR = "\n" | |
| class LocalChatbot: | |
| def __init__(self, llm, memory, tokenizer=None, default_template: Optional[str] = "general"): | |
| self.llm = llm | |
| self.memory = memory | |
| self.tokenizer = tokenizer | |
| self.default_template = default_template | |
| def _build_system_prompt(self, intent: str) -> str: | |
| return TEMPLATES.get(intent, TEMPLATES.get(self.default_template, TEMPLATES["general"])) | |
| def _build_prompt(self, user_message: str, intent: str, max_pairs: int = 12) -> str: | |
| try: | |
| self.memory.trim_to_recent_pairs(max_pairs) | |
| except Exception: | |
| pass | |
| system_prompt = self._build_system_prompt(intent) | |
| history_text = self.memory.get_formatted(separator=MSG_SEPARATOR) | |
| parts = [ | |
| f"System: {system_prompt}", | |
| history_text, | |
| f"User: {user_message}", | |
| "Assistant:" | |
| ] | |
| return MSG_SEPARATOR.join([p for p in parts if p]) | |
| def ask(self, user_message: str, gen_args: Optional[Dict[str, Any]] = None) -> str: | |
| if not user_message.strip(): | |
| return "Please enter a message." | |
| intent = detect_intent(user_message) | |
| prompt = self._build_prompt(user_message, intent) | |
| gen = DEFAULT_GEN_ARGS.copy() | |
| if gen_args: | |
| gen.update(gen_args) | |
| try: | |
| if self.tokenizer: | |
| # Transformers-style generation | |
| inputs = self.tokenizer(prompt, return_tensors="pt") | |
| outputs = self.llm.generate(**inputs, max_new_tokens=gen.get("max_tokens", 300)) | |
| bot_reply = self.tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| else: | |
| # Fallback: callable LLM | |
| bot_reply = self.llm(prompt, **gen) | |
| if isinstance(bot_reply, dict) and "choices" in bot_reply: | |
| bot_reply = bot_reply["choices"][0].get("text", "").strip() | |
| except Exception: | |
| bot_reply = "Sorry — I couldn't generate a response. Please try again." | |
| if not bot_reply: | |
| bot_reply = "Sorry — I couldn't generate a response. Please try again." | |
| try: | |
| self.memory.add(user_message, bot_reply) | |
| except Exception: | |
| try: | |
| self.memory.add_message("user", user_message) | |
| self.memory.add_message("assistant", bot_reply) | |
| except Exception: | |
| pass | |
| return bot_reply | |