File size: 2,942 Bytes
3740bd7 61d9448 3740bd7 61d9448 3740bd7 | 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | """
Response Generation module for XENO Bot
Handles LLM response generation
"""
from typing import Dict, List
from transformers import pipeline
from src.config import (HF_TOKEN, LLM_MAX_NEW_TOKENS, LLM_MODEL_NAME,
LLM_TEMPERATURE, SYSTEM_PROMPT)
_text_generator = None
def get_text_generator():
"""Lazily load and cache the local Transformers text-generation pipeline."""
global _text_generator
if _text_generator is None:
_text_generator = pipeline(
"text-generation",
model=LLM_MODEL_NAME,
token=HF_TOKEN,
device_map="auto",
)
return _text_generator
def _call_local_llm(prompt: str) -> str:
"""Generate text locally using the Transformers pipeline."""
generator = get_text_generator()
outputs = generator(
prompt,
max_new_tokens=LLM_MAX_NEW_TOKENS,
do_sample=True,
temperature=LLM_TEMPERATURE,
return_full_text=False,
)
text = outputs[0].get("generated_text", "") if outputs else ""
if not text:
raise ValueError("Local model returned an empty response")
return text.strip()
def generate_xeno_response(
context: str, question: str, chat_history: List[Dict[str, str]], timer=None
) -> str:
"""
Generate a response using the LLM
Args:
context: Formatted context from knowledge base
question: User's question
chat_history: List of previous messages
timer: Optional timer object for tracking
Returns:
Generated response text
"""
if timer:
with timer.time_step("llm_generation"):
return _generate_response_impl(context, question, chat_history)
else:
return _generate_response_impl(context, question, chat_history)
def _generate_response_impl(
context: str, question: str, chat_history: List[Dict[str, str]]
) -> str:
"""Internal implementation of response generation"""
# Format chat history
formatted_history = (
"\n".join(
[f"{msg['role'].capitalize()}: {msg['content']}" for msg in chat_history]
)
if chat_history
else "None"
)
# Build prompt
prompt = f"{SYSTEM_PROMPT}\n### HISTORY ###\n{formatted_history}\n### CONTEXT ###\n{context}\n### QUESTION ###\n{question}"
return _call_local_llm(prompt)
def format_chat_history(messages: List[Dict[str, str]]) -> str:
"""
Format chat history for display or logging
Args:
messages: List of message dictionaries with 'role' and 'content'
Returns:
Formatted string representation of chat history
"""
if not messages:
return "No previous conversation"
formatted = []
for msg in messages:
role = msg.get("role", "unknown").capitalize()
content = msg.get("content", "")
formatted.append(f"{role}: {content}")
return "\n".join(formatted)
|