app.py

import streamlit as st
from datasets import load_dataset
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel, pipeline
from sentence_transformers import SentenceTransformer
import time
from datetime import datetime
import json
import os
import pickle
import random
import re
import warnings
warnings.filterwarnings('ignore')

# ===================================================================
# 1. НАСТРОЙКИ
# ===================================================================

print(f"CUDA доступна: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")

MODEL_NAME = "ai-forever/rugpt3small_based_on_gpt2"
EMBEDDING_MODEL = "all-MiniLM-L6-v2"
SCIENCE_DATASET = "RafaelUI/ru_science"
ARTICLE_LIMIT = 50
MAX_LENGTH = 512
TEMPERATURE = 0.8
TOP_P = 0.95

LOG_FILE = "query_logs.json"
EMBEDDINGS_FILE = "science_embeddings.npy"
ARTICLES_FILE = "science_articles.pkl"

AI_NAME = "OpenAirAI"
COMPANY_NAME = "OpenRussianAI"
CREATORS = ["Грибков Евгений", "RootLinux21"]
WEBSITE = "https://sites.google.com/view/opruai/home"
HUGGINGFACE = "https://huggingface.co/OpenRussianAI"
CREATION_DATE = "2026"

st.set_page_config(
    page_title=f"{AI_NAME} - Научный AI-ассистент",
    page_icon="🧠",
    layout="wide",
    initial_sidebar_state="expanded"
)

# ===================================================================
# 2. НАСТОЯЩАЯ НЕЙРОСЕТЬ (ВСЕГДА ГЕНЕРИРУЕТ)
# ===================================================================

class NeuralChatbot:
    def __init__(self):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.tokenizer = None
        self.model = None
        self.generator = None
        self.is_loaded = False
        
        # Системный промпт для нейросети
        self.system_prompt = f"""Ты - {AI_NAME}, дружелюбный научный AI-ассистент от компании {COMPANY_NAME}.
Ты создан в {CREATION_DATE} командой {', '.join(CREATORS)}.
Ты всегда отвечаешь на русском языке, тепло и профессионально.
Ты помогаешь с научными вопросами, системным администрированием и технологиями.
Ты генерируешь уникальные ответы каждый раз.

Вот вопрос пользователя: """
    
    def load_model(self):
        with st.spinner("🧠 Загружаю нейросеть..."):
            try:
                from transformers import GPT2LMHeadModel, GPT2Tokenizer
                
                self.tokenizer = GPT2Tokenizer.from_pretrained(MODEL_NAME)
                self.model = GPT2LMHeadModel.from_pretrained(MODEL_NAME)
                self.model.to(self.device)
                self.model.eval()
                
                self.generator = pipeline(
                    'text-generation',
                    model=self.model,
                    tokenizer=self.tokenizer,
                    device=0 if torch.cuda.is_available() else -1,
                    max_length=300,
                    temperature=TEMPERATURE,
                    top_p=TOP_P,
                    do_sample=True,
                    repetition_penalty=1.2,
                    pad_token_id=self.tokenizer.eos_token_id
                )
                
                self.is_loaded = True
                return True
            except Exception as e:
                st.warning(f"Не удалось загрузить нейросеть: {e}")
                return False
    
    def generate(self, query):
        """Генерация ответа нейросетью"""
        if not self.is_loaded:
            return self.fallback_response(query)
        
        try:
            # Формируем промпт
            prompt = self.system_prompt + query
            
            # Генерируем
            response = self.generator(
                prompt,
                max_new_tokens=250,
                temperature=TEMPERATURE,
                top_p=TOP_P,
                do_sample=True,
                repetition_penalty=1.2
            )[0]['generated_text']
            
            # Убираем промпт
            response = response.replace(prompt, "").strip()
            
            # Если ответ пустой или слишком короткий
            if len(response) < 15:
                return self.fallback_response(query)
            
            return response
            
        except Exception as e:
            print(f"Ошибка генерации: {e}")
            return self.fallback_response(query)
    
    def fallback_response(self, query):
        """Резервный ответ (если нейросеть не работает)"""
        return f"""Я {AI_NAME} от {COMPANY_NAME}. 

К сожалению, сейчас нейросеть временно недоступна, но я хочу ответить на ваш вопрос: "{query}"

Обычно я генерирую уникальные ответы с помощью нейросети GPT. Попробуйте позже, и я смогу дать более развернутый ответ! 🧠

А пока я могу рассказать, что создан в {CREATION_DATE} командой {', '.join(CREATORS)}. Я помогаю с научными вопросами и технологиями."""

# ===================================================================
# 3. ЗАГРУЗКА СТАТЕЙ (ДЛЯ КОНТЕКСТА)
# ===================================================================

@st.cache_resource
def load_science_articles():
    if os.path.exists(ARTICLES_FILE):
        with open(ARTICLES_FILE, 'rb') as f:
            return pickle.load(f)
    
    with st.spinner("📚 Загружаю научные статьи..."):
        try:
            dataset = load_dataset(SCIENCE_DATASET, split="train", streaming=True)
            articles = []
            for i, row in enumerate(dataset):
                if i >= ARTICLE_LIMIT:
                    break
                text = row.get('content', '') or row.get('text', '') or str(row)
                title = row.get('title', f"Статья {i}")
                articles.append({
                    "id": i,
                    "title": title[:200],
                    "text": text[:2000],
                    "source": "ru_science"
                })
            with open(ARTICLES_FILE, 'wb') as f:
                pickle.dump(articles, f)
            return articles
        except Exception as e:
            st.warning(f"Не удалось загрузить статьи: {e}")
            return []

@st.cache_resource
def load_embedder():
    return SentenceTransformer(EMBEDDING_MODEL)

@st.cache_resource
def create_embeddings(_articles, _embedder):
    if os.path.exists(EMBEDDINGS_FILE):
        return np.load(EMBEDDINGS_FILE)
    if not _articles:
        return np.array([])
    texts = [f"{a['title']}\n\n{a['text']}" for a in _articles]
    embeddings = _embedder.encode(texts, normalize_embeddings=True, show_progress_bar=True, batch_size=64)
    np.save(EMBEDDINGS_FILE, embeddings)
    return embeddings

def search_articles(query, _articles, _embeddings, _embedder):
    """Поиск релевантных статей"""
    if not _articles or len(_embeddings) == 0:
        return []
    query_vector = _embedder.encode([query], normalize_embeddings=True)[0]
    scores = _embeddings @ query_vector
    top_indices = np.argsort(-scores)[:2]
    results = []
    for idx in top_indices:
        score = float(scores[int(idx)])
        if score > 0.15:
            article = _articles[int(idx)]
            results.append({"title": article['title'], "score": score, "text": article['text'][:500]})
    return results

# ===================================================================
# 4. ОЧИСТКА ЗАПРОСОВ
# ===================================================================

def clean_query(query):
    """Очищает запрос от спама"""
    query = re.sub(r'http[s]?://\S+', '', query)
    query = re.sub(r'\S+@\S+', '', query)
    query = re.sub(r'\+7\s*\(?\d{3}\)?\s*\d{3}\s*\d{2}\s*\d{2}', '', query)
    query = re.sub(r'[^\w\s\.\?\!,;:]', ' ', query)
    query = ' '.join(query.split())
    
    if len(query) > 300:
        sentences = re.split(r'[.!?]', query)
        query = '. '.join(sentences[:2]) + '.' if sentences else query[:200]
    
    return query.strip()

def enhance_with_context(query, articles_context):
    """Добавляет контекст из статей к запросу"""
    if not articles_context:
        return query
    
    context = "\n\nВот релевантная научная информация:\n"
    for i, art in enumerate(articles_context, 1):
        context += f"{i}. {art['title']}\n{art['text'][:300]}...\n"
    
    context += f"\nНа основе этой информации, ответь на вопрос: {query}"
    return context

# ===================================================================
# 5. ОСНОВНОЙ КЛАСС
# ===================================================================

class OpenAirAI:
    def __init__(self):
        self.name = AI_NAME
        self.company = COMPANY_NAME
        self.creators = CREATORS
        self.chatbot = NeuralChatbot()
        self.is_ready = False
        
    def initialize(self):
        self.is_ready = self.chatbot.load_model()
        return self.is_ready
    
    def generate_answer(self, query, articles_context=None):
        """Генерирует ответ ТОЛЬКО нейросетью, без if/else"""
        clean_q = clean_query(query)
        
        # Если есть контекст статей - добавляем его
        if articles_context:
            enhanced_query = enhance_with_context(clean_q, articles_context)
        else:
            enhanced_query = clean_q
        
        # ВСЕГДА генерируем нейросетью
        return self.chatbot.generate(enhanced_query)

# ===================================================================
# 6. ИНТЕРФЕЙС
# ===================================================================

# Загрузка данных
articles = load_science_articles()
embedder = load_embedder()
embeddings = create_embeddings(articles, embedder)

# Инициализация AI
if 'ai' not in st.session_state:
    st.session_state.ai = OpenAirAI()
    st.session_state.ai.initialize()

ai = st.session_state.ai

# История чата
if "messages" not in st.session_state:
    st.session_state.messages = []
    # Первое приветствие генерируется нейросетью
    greeting = ai.generate_answer("Привет! Представься и расскажи о себе")
    st.session_state.messages.append({"role": "assistant", "content": greeting})

# --- БОКОВАЯ ПАНЕЛЬ ---
with st.sidebar:
    st.image("https://cdn-icons-png.flaticon.com/512/4248/4248455.png", width=80)
    st.title(f"🧠 {AI_NAME}")
    
    st.markdown(f"""
    **{ai.name}** | {CREATION_DATE}
    
    **Компания:** {ai.company}
    **Разработчики:** {', '.join(ai.creators)}
    
    ---
    **🔗 Ссылки:**
    [🌐 Сайт]({WEBSITE})
    [🤗 Hugging Face]({HUGGINGFACE})
    
    ---
    **📊 Статистика:**
    - Статей: {len(articles)}
    - Сообщений: {len(st.session_state.messages)}
    - Модель: {MODEL_NAME.split("/")[-1]}
    """)
    
    if st.button("🗑️ Очистить чат"):
        st.session_state.messages = []
        greeting = ai.generate_answer("Привет! Представься и расскажи о себе")
        st.session_state.messages.append({"role": "assistant", "content": greeting})
        st.rerun()
    
    if st.button("🔄 Сгенерировать новое приветствие"):
        new_greeting = ai.generate_answer("Привет! Представься заново")
        if st.session_state.messages:
            st.session_state.messages[0] = {"role": "assistant", "content": new_greeting}
        st.rerun()

# --- ОСНОВНАЯ ЧАСТЬ ---
st.title(f"🧠 {AI_NAME} - Нейросетевой AI-ассистент")
st.markdown(f"**{AI_NAME}** от **{COMPANY_NAME}** | Ответы генерируются нейросетью")

# Отображение сообщений
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Поле ввода
if prompt := st.chat_input("Задайте вопрос..."):
    # Добавляем сообщение пользователя
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)
    
    # Генерация ответа нейросетью
    with st.chat_message("assistant"):
        with st.spinner("🧠 Нейросеть генерирует ответ..."):
            # Ищем релевантные статьи для контекста
            articles_context = search_articles(prompt, articles, embeddings, embedder)
            
            # Генерируем ответ (ВСЕГДА нейросетью)
            response = ai.generate_answer(prompt, articles_context)
            
            # Добавляем статьи в ответ, если они есть и не встроены
            if articles_context and len(response) < 50:
                response += "\n\n📄 Я нашел релевантные научные статьи:\n"
                for i, art in enumerate(articles_context, 1):
                    response += f"{i}. {art['title']}\n"
            
            st.markdown(response)
            st.session_state.messages.append({"role": "assistant", "content": response})
    
    st.rerun()

# --- ПОДВАЛ ---
st.divider()
st.caption(f"🧠 {AI_NAME} от {COMPANY_NAME} | Создан в {CREATION_DATE} | Генеративная нейросеть")