test5 / chatbot.py
vydrking's picture
Upload 18 files
2fc8dc5 verified
import re
import json
from typing import List, Dict, Tuple
from knowledge_base import KnowledgeBase
from retriever import Retriever
class ITMOChatbot:
def __init__(self):
self.kb = KnowledgeBase()
self.retriever = Retriever()
self.max_history_turns = 3
self.max_context_tokens = 1200
self.relevance_threshold = 0.38
try:
from transformers import pipeline
self.generator = pipeline('text2text-generation', model='cointegrated/rut5-base-multitask')
except Exception as e:
print(f'Генеративная модель не загружена: {e}')
self.generator = None
def chat(self, message: str, history: list) -> Tuple[str, float]:
if not message.strip():
return 'Пожалуйста, задайте вопрос.', 0.0
if not self.kb.is_itmo_query(message):
return self._get_irrelevant_response(), 0.0
context = self._get_context(message)
if not context:
return 'К сожалению, не нашел релевантной информации в учебных планах ITMO.', 0.0
response = self._generate_response(message, history, context)
relevance_score = self._calculate_relevance_score(message, context)
return response, relevance_score
def recommend_courses(self, profile: dict) -> str:
if not profile.get('semester'):
return 'Пожалуйста, укажите целевой семестр для получения рекомендаций.'
recommendations = self.kb.recommend(profile)
if not recommendations:
return 'К сожалению, не удалось найти подходящие курсы для вашего профиля.'
result = '🎯 Рекомендуемые курсы (из официальных учебных планов ITMO):\n\n'
for i, rec in enumerate(recommendations[:7], 1):
result += f'{i}. {rec["name"]} ({rec["semester"]} семестр, {rec["credits"]} кредитов)\n'
result += f' {rec["why"]}\n\n'
return result
def _get_context(self, message: str) -> List[Dict]:
try:
results = self.retriever.retrieve(message, k=6, threshold=0.35)
# Преобразуем результаты в нужный формат
formatted_results = []
for result in results:
course_id = result.get('course_id')
if course_id:
course = self.kb.get_course_by_id(course_id)
if course:
course['score'] = result.get('score', 0.0)
formatted_results.append(course)
return formatted_results
except Exception as e:
print(f'Ошибка при получении контекста: {e}')
return []
def _generate_response(self, message: str, history: list, context: List[Dict]) -> str:
if not context:
return 'В предоставленных данных об этом не сказано.'
prompt = self._build_prompt(message, history, context)
if self.generator:
try:
response = self.generator(
prompt,
max_new_tokens=180,
temperature=0.4,
do_sample=True
)[0]['generated_text']
return response.strip()
except Exception as e:
print(f'Ошибка генерации: {e}')
return self._fallback_response(context)
def _build_prompt(self, message: str, history: list, context: List[Dict]) -> str:
system_prompt = 'Отвечай только по контексту (ниже). Если недостаточно данных — прямо скажи: "в предоставленных данных об этом не сказано". Отвечай кратко и по делу.'
history_text = ''
if history:
recent_history = history[-self.max_history_turns:]
for turn in recent_history:
history_text += f'Пользователь: {turn[0]}\nБот: {turn[1]}\n'
context_text = 'Контекст:\n'
for item in context:
context_text += f'- {item["name"]} ({item["semester"]} семестр, {item["credits"]} кредитов): {item["short_desc"]}\n'
prompt = f'{system_prompt}\n\n{history_text}Контекст:\n{context_text}\nВопрос: {message}'
if len(prompt) > self.max_context_tokens * 4:
prompt = prompt[:self.max_context_tokens * 4]
return prompt
def _fallback_response(self, context: List[Dict]) -> str:
if not context:
return 'В предоставленных данных об этом не сказано.'
courses = []
for item in context[:3]:
courses.append(f'{item["name"]} ({item["semester"]} семестр, {item["credits"]} кредитов)')
return f'Найденные курсы: {", ".join(courses)}. Для более подробной информации обратитесь к официальным учебным планам ITMO.'
def _calculate_relevance_score(self, message: str, context: List[Dict]) -> float:
if not context:
return 0.0
scores = [item.get('score', 0.0) for item in context]
return sum(scores) / len(scores) if scores else 0.0
def _get_irrelevant_response(self) -> str:
return '''Похоже, вопрос не относится к магистратурам ITMO и их учебным планам.
Попробуйте спросить, например:
• "Какие дисциплины по NLP в 1 семестре программы ИИ?"
• "Расскажи о программе AI Product"
• "Какие курсы по машинному обучению есть в программе ИИ?"
• "Сколько кредитов за дисциплину 'Глубокое обучение'?"'''