import os
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# Токен, если модель приватная
hf_token = os.getenv("HF_TOKEN", None)
model_path = "./capybara-finetuned"  # или HF-репозиторий, например: "NousResearch/Nous-Capybara-3B-V1.9"

# Загружаем модель и токенизатор
tokenizer = AutoTokenizer.from_pretrained(
    model_path,
    token=hf_token,
    trust_remote_code=True,
    use_fast=True,
)

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    token=hf_token,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto",
    trust_remote_code=True,
)

# Создаём пайплайн
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Функция классификации
def classify(text):
    prompt = f"### Вопрос:\n{text}\n\n### Класс:"
    try:
        result = pipe(prompt, max_new_tokens=10, do_sample=False)[0]["generated_text"]
        label = result.split("### Класс:")[-1].strip().split()[0].lower()
        return f"🔍 Класс: **{label}**"
    except Exception as e:
        return f"❌ Ошибка: {str(e)}"

# Интерфейс Gradio
iface = gr.Interface(
    fn=classify,
    inputs=gr.Textbox(lines=4, placeholder="Введите сообщение..."),
    outputs="markdown",
    title="Capybara Text Classifier 🦫",
    description="Классификация текста как 'запрос' или 'реклама' с помощью Capybara-3B",
    # enable_api=True,           # Разрешаем вызывать данный Interface извне
    # api_name="/classify"       # Название эндпоинта (путь для client.predict)
)

app, local_url, share_url = iface.launch(share=True, ssr_mode=False)