Spaces:
Sleeping
Sleeping
File size: 2,702 Bytes
0885208 15fa2e5 0885208 15fa2e5 0885208 15fa2e5 0885208 15fa2e5 0885208 15fa2e5 0885208 15fa2e5 0885208 15fa2e5 0885208 15fa2e5 0885208 15fa2e5 0885208 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
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
|