app.py

import os
import random
import time
import shutil
from datetime import datetime, timezone
import pandas as pd
import gradio as gr
from datasets import load_dataset, Dataset
from transformers import (
    AutoTokenizer, 
    GPT2LMHeadModel, 
    GPT2Config, 
    TrainingArguments, 
    Trainer,
    DataCollatorForLanguageModeling
)
from huggingface_hub import model_info

# --- CSS DO UKRYCIA NIECHCIANYCH ELEMENTÓW ---
custom_css = """
.header-bar, div[data-testid="header"], .app-header, .flex-row.items-center.justify-between {
    display: none !important; height: 0 !important;
}
footer, .built-with-gradio, footer.svelte-1y4jv3v, div.svelte-1y4jv3v, .settings-btn, 
button[aria-label="Settings"], .gradio-container > footer, footer > *, footer a, footer button {
    display: none !important; visibility: hidden !important; height: 0 !important;
    opacity: 0 !important; pointer-events: none !important;
}
"""

# --- USTAWIENIA ---
hf_token = os.environ.get("HF_TOKEN")
dataset_name = "StrawberryJelly/unimind-kronika"
model_repo = "StrawberryJelly/unimind-mozg"
model_dir = "./model"

# --- HOT RELOAD ZMIENNE ---
last_model_check = time.time()
last_model_version = None

# --- INIT KRONIKI ---
try:
    kronika = load_dataset(dataset_name, split="train", token=hf_token)
    if len(kronika) == 0: raise ValueError("Empty dataset")
except Exception:
    now_utc = datetime.now(timezone.utc).isoformat()
    init_data = [{"id": 1, "ts": now_utc, "user": "StrawberryJelly", "q": "genesis", "a": "a", "flag": "ok"}]
    kronika = Dataset.from_list(init_data)
    kronika.push_to_hub(dataset_name, token=hf_token)

# --- INIT MODELU ---
os.makedirs(model_dir, exist_ok=True)
tokenizer = AutoTokenizer.from_pretrained("flax-community/papuGaPT2")
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

try:
    model = GPT2LMHeadModel.from_pretrained(model_repo, token=hf_token)
    print("✅ Załadowano istniejący model z HF Hub.")
except Exception:
    print("️ Brak modelu na Hubie. Tworzę losowy model od zera...")
    config = GPT2Config(vocab_size=50257, n_positions=1024, n_embd=768, n_layer=12, n_head=12)
    model = GPT2LMHeadModel(config)
    model.save_pretrained(model_dir)
    tokenizer.save_pretrained(model_dir)
    model.push_to_hub(model_repo, token=hf_token)
    tokenizer.push_to_hub(model_repo, token=hf_token)

# --- CZYSZCZENIE CACHE I HOT RELOAD ---
def clean_model_cache():
    """Czyści lokalne pliki modelu, żeby nie zapychać dysku"""
    try:
        if os.path.exists("./tmp"):
            shutil.rmtree("./tmp")
        if os.path.exists("./model"):
            shutil.rmtree("./model")
            os.makedirs("./model", exist_ok=True)
        print("🧹 Lokalny folder modelu wyczyszczony.")
    except Exception as e:
        print(f"⚠️ Błąd czyszczenia folderu: {e}")

def check_for_new_model():
    """Sprawdź co 10 minut czy jest nowy model na HF Hub"""
    global model, last_model_check, last_model_version
    current_time = time.time()
    
    if current_time - last_model_check < 600:  # Sprawdź co 10 minut
        return
    
    last_model_check = current_time
    
    try:
        info = model_info(model_repo, token=hf_token)
        current_version = info.lastModified
        
        if last_model_version is None:
            last_model_version = current_version
            print(f"📦 Wersja modelu: {current_version}")
        elif current_version != last_model_version:
            print(f"🔄 Nowa wersja modelu wykryta! Sprzątam i przeładowuję...")
            clean_model_cache()
            model = GPT2LMHeadModel.from_pretrained(model_repo, token=hf_token)
            last_model_version = current_version
            print(f"✅ Model przeładowany!")
    except Exception as e:
        print(f"⚠️ Błąd hot-reload: {e}")

# --- STATUS I MATEMATYKA ---
def get_status():
    global kronika
    try:
        ds = load_dataset(dataset_name, split="train", token=hf_token)
    except Exception:
        ds = kronika
    total = len(ds)
    ok_count = len(ds.filter(lambda x: x['flag'] == 'ok'))
    pending_count = len(ds.filter(lambda x: x['flag'] == 'pending'))
    return f"UNIMIND. Stan: Przetrenowane: {ok_count} | Oczekujące: {pending_count} | Łącznie: {total}"

def generate_math():
    a = random.randint(2, 9)
    b = random.randint(2, 9)
    c = random.randint(1, 5)
    if random.choice([True, False]):
        question = f"{a} * {b} + {c} = ?"
        ans = a * b + c
    else:
        question = f"{a} * {b} - {c} = ?"
        ans = a * b - c
    return question, str(ans)

# --- JAVASCRIPT DO LOCALSTORAGE ---
js_save_chat = """
function(chat_history) { 
    if (!Array.isArray(chat_history)) return [];
    localStorage.setItem('unimind_chat_history', JSON.stringify(chat_history)); 
    return chat_history; 
}"""

js_load_chat = """
function() { 
    const saved = localStorage.getItem('unimind_chat_history'); 
    if (saved) { 
        try { 
            const parsed = JSON.parse(saved);
            return Array.isArray(parsed) ? parsed.filter(m => typeof m === 'object' && m.role && m.content) : [];
        } catch(e) { return []; } 
    } 
    return []; 
}"""

js_clear_chat = "function() { localStorage.removeItem('unimind_chat_history'); return []; }"

# --- UI ---
# POPRAWKA 1: Usunięto css z Blocks()
with gr.Blocks() as demo:
    math_q, math_a = generate_math()
    math_state = gr.State(math_a)
    math_q_display = gr.Markdown(f"### Zabezpieczenie antybot: `{math_q}`")
    
    ans_input = gr.Textbox(label="Podaj wynik")
    submit_btn = gr.Button("Zatwierdź")
    
    app_container = gr.Column(visible=False)
    
    def check_ans(user_ans, correct_ans):
        if str(user_ans).strip() == str(correct_ans).strip():
            return gr.update(visible=False), gr.update(visible=True), correct_ans, ""
        else:
            q, a = generate_math()
            return gr.update(value=f"### Zabezpieczenie antybot: `{q}` \n\n **Błędna odpowiedź.**"), gr.update(visible=False), a, ""

    submit_btn.click(check_ans, [ans_input, math_state], [math_q_display, app_container, math_state, ans_input])
    
    with app_container:
        status_md = gr.Markdown(get_status())
        
        with gr.Tabs():
            with gr.TabItem("Gadaj"):
                gr.Warning("⚠️ Trening modelu odbywa się o 3:00, 12:00 i 22:00 czasu polskiego. W tych godzinach serwer może działać wolniej.")
                gr.Markdown("UWAGA: Na start będzie bełkotał losowe tokeny. To normalne.")
                
                # POPRAWKA 2: Usunięto type="messages" i show_copy_button (domyślne w Gradio 6.0)
                chatbot = gr.Chatbot(value=[], height=400)
                msg = gr.Textbox(label="Twój prompt")
                clear_btn = gr.Button("🗑️ Wyczyść rozmowę")
                
                def respond(message, chat_history):
                    check_for_new_model()  # Hot reload check
                    
                    if chat_history is None:
                        chat_history = []
                    if not message or not message.strip():
                        return chat_history
                    
                    context = ""
                    if chat_history:
                        context = "\n--- HISTORIA ---\n"
                        for msg_pair in chat_history[-2:]:
                            if isinstance(msg_pair, dict):
                                role = msg_pair.get('role', '')
                                content = msg_pair.get('content', '')
                                if role == 'user': context += f"Pytanie: {content}\n"
                                elif role == 'assistant': context += f"Odpowiedź: {content}\n"
                        context += "--- KONIEC HISTORII ---\n\n"
                    
                    formatted = f"{context}NOWE PYTANIE: {message}\nOdpowiedź:"
                    inputs = tokenizer(formatted, return_tensors="pt", truncation=True, max_length=512)
                    
                    outputs = model.generate(
                        **inputs, max_new_tokens=50, temperature=1.0, do_sample=True,
                        repetition_penalty=1.2, pad_token_id=tokenizer.eos_token_id
                    )
                    
                    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
                    
                    if "Odpowiedź:" in decoded:
                        parts = decoded.split("Odpowiedź:")
                        ans = parts[-1].strip()
                    else:
                        ans = decoded.strip()

                    chat_history.append({"role": "user", "content": message})
                    chat_history.append({"role": "assistant", "content": ans})
                    return chat_history

                msg.submit(respond, [msg, chatbot], [chatbot]).then(None, [chatbot], None, js=js_save_chat)
                clear_btn.click(None, None, chatbot, js=js_clear_chat)
                
            with gr.TabItem("Naucz"):
                gr.Markdown("### 📝 Dodaj wpis do Kroniki")
                gr.Markdown("Wpis trafi do bazy z flagą 'pending'. Model nauczy się go podczas najbliższego treningu (3:00, 12:00, 22:00 PL).")
                
                q_in = gr.Textbox(label="Kiedy ktoś zapyta:")
                a_in = gr.Textbox(label="Unimind ma odpowiedzieć:")
                add_btn = gr.Button("➕ DODAJ DO KRONIKI")
                add_status = gr.Markdown("")
                
                def add_to_kronika(q, a):
                    global kronika
                    try:
                        kronika = load_dataset(dataset_name, split="train", token=hf_token)
                    except Exception:
                        pass
                        
                    new_id = len(kronika) + 1
                    now_utc = datetime.now(timezone.utc).isoformat()
                    new_row = {"id": new_id, "ts": now_utc, "user": "anon", "q": q, "a": a, "flag": "pending"}
                    
                    df = kronika.to_pandas()
                    df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
                    kronika = Dataset.from_pandas(df)
                    kronika.push_to_hub(dataset_name, token=hf_token)
                    
                    pending_count = len(kronika.filter(lambda x: x['flag'] == 'pending'))
                    total_count = len(kronika)
                    
                    return f"✅ Dodano! Oczekujących wpisów: {pending_count}/{total_count}.", get_status()

                add_btn.click(add_to_kronika, [q_in, a_in], [add_status, status_md])
                
            with gr.TabItem("Kronika"):
                gr.Markdown(f"[Otwórz Kronikę na Hugging Face](https://huggingface.co/datasets/{dataset_name})")
                gr.Markdown("Zmień flag na 'rejected' żeby usunąć neuron z treningu")

    demo.load(None, None, chatbot, js=js_load_chat)

# POPRAWKA 3: Przeniesiono css do launch()
demo.launch(css=custom_css)