|
|
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"
|
|
|
• "Какие курсы по машинному обучению есть в программе ИИ?"
|
|
|
• "Сколько кредитов за дисциплину 'Глубокое обучение'?"'''
|
|
|
|