app.py

# app.py
import os
import json
import re
import torch
from difflib import get_close_matches
from transformers import AutoTokenizer, AutoModelForCausalLM
from ddgs import DDGS
import gradio as gr

# === Загрузка patterns.json ===
PATTERNS = {}
HAS_PATTERNS = False
try:
    with open("patterns.json", "r", encoding="utf-8") as f:
        PATTERNS = json.load(f)
    HAS_PATTERNS = True
except:
    pass

KEYWORDS = {
    "привет": ["привет", "здравствуй", "хай"],
    "как дела": ["дела", "как ты", "настроение"],
    "имя": ["имя", "кто ты", "зовут"],
    "пока": ["пока", "выход", "до свидания", "стоп"]
}

def preprocess(text):
    return re.sub(r'[^а-яё\s]', ' ', text.lower()).strip()

# === Загрузка TinyLlama (без токена!) ===
print("Загрузка TinyLlama с Hugging Face Hub...")
MODEL_ID = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=False)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto",
    low_cpu_mem_usage=True
)
print("✅ TinyLlama готова!")

# === Веб-поиск ===
def web_search(query, max_results=3):
    try:
        with DDGS() as ddgs:
            results = ddgs.text(query, region="ru-ru", max_results=max_results)
            return "\n".join([f"{r['title']}: {r['body']}" for r in results])
    except Exception as e:
        return f"Ошибка поиска: {str(e)[:100]}"

# === Генерация через TinyLlama ===
def generate_with_tinyllama(prompt_text, max_tokens=256):
    try:
        inputs = tokenizer(prompt_text, return_tensors="pt").to(model.device)
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            do_sample=True,
            temperature=0.6,
            top_p=0.92,
            pad_token_id=tokenizer.eos_token_id,
            repetition_penalty=1.1
        )
        full = tokenizer.decode(outputs[0], skip_special_tokens=True)
        if "<|assistant|>" in full:
            return full.split("<|assistant|>")[-1].strip()
        else:
            return full[len(prompt_text):].strip()
    except Exception as e:
        return f"Ошибка генерации: {str(e)[:100]}"

# === Fallback через patterns.json ===
def get_fallback_response(user_input):
    if not HAS_PATTERNS:
        return None
    clean = preprocess(user_input)
    if not clean:
        return None
    if any(w in clean for w in ["пока", "выход", "стоп", "до свидания"]):
        return random.choice(PATTERNS["пока"])
    knowledge = PATTERNS.get("knowledge", {})
    if clean in knowledge:
        return knowledge[clean]
    matches = get_close_matches(clean, knowledge.keys(), n=1, cutoff=0.6)
    if matches:
        return knowledge[matches[0]]
    for intent, words in KEYWORDS.items():
        if any(get_close_matches(token, words, n=1, cutoff=0.6) for token in clean.split()):
            return random.choice(PATTERNS[intent])
    return None

# === Основная логика ===
def respond(message, history):
    user_input = message.strip()
    user_lower = user_input.lower()

    # === Системные команды (заглушка) ===
    if user_lower.startswith("система:"):
        return "🔒 Управление ОС недоступно в демо."

    # === Перевод ===
    if user_lower.startswith("перевод:"):
        text = user_input[8:].strip()
        if not text:
            return "🔤 Пример: `перевод: Hello, how are you?`"
        prompt = f"<|user|>\nПереведи на русский язык: {text}\n<|assistant|>\n"
        return generate_with_tinyllama(prompt, max_tokens=128)

    # === Веб-поиск ===
    if user_lower.startswith("поиск:"):
        query = user_input[6:].strip()
        if not query:
            return "🔍 Пример: `поиск: погода в Москве`"
        context = web_search(query)
        prompt = (
            f"<|user|>\nИнформация из интернета:\n{context}\n\n"
            f"Кратко ответь на русском: {query}\n"
            f"<|assistant|>\n"
        )
        return generate_with_tinyllama(prompt, max_tokens=256)

    # === Fallback через patterns.json ===
    fallback = get_fallback_response(user_input)
    if fallback:
        return fallback

    # === Обычный режим ===
    prompt = f"<|user|>\n{user_input}\nОтветь кратко на русском.\n<|assistant|>\n"
    return generate_with_tinyllama(prompt, max_tokens=256)

# === Gradio интерфейс ===
chatbot = gr.ChatInterface(
    respond,
    title="🚀 Newton MAX (TinyLlama + поиск + перевод)",
    description="Поддержка команд: `поиск: ...`, `перевод: ...`, `система: ...`",
    examples=[
        "Привет!",
        "Что такое квантовый компьютер?",
        "поиск: курс доллара",
        "перевод: How are you doing today?"
    ],
    theme="soft"
)

if __name__ == "__main__":
    chatbot.launch()