# app.py
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F

# 1. Tải mô hình và tokenizer từ Hugging Face Hub
model_name = "MonkeyDAnh/Model-AI-Percentage-used-v2" # Thay thế bằng tên repo model của bạn
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Kiểm tra nếu có GPU thì sử dụng GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval() # Chuyển mô hình sang chế độ đánh giá

# 2. Định nghĩa hàm dự đoán
def predict_ai_human(text):
    """
    Hàm dự đoán tỷ lệ AI/Human dựa trên văn bản đầu vào.
    """
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
    inputs = {k: v.to(device) for k, v in inputs.items()} # Chuyển inputs sang device

    with torch.no_grad():
        outputs = model(**inputs)

    # Lấy logits và chuyển đổi sang xác suất
    # Giả sử mô hình của bạn có 2 lớp: 0 (Human), 1 (AI)
    # Hoặc ngược lại, tùy thuộc vào cách bạn huấn luyện
    # Cần kiểm tra thứ tự nhãn của mô hình của bạn!
    # Ví dụ: model.config.id2label nếu có
    # Đối với model binary classification, thường có 2 output logits
    # Output thường là logits cho mỗi lớp.
    # Nếu mô hình của bạn dự đoán "AI percentage", thì output có thể là 1 giá trị duy nhất
    # hoặc 2 giá trị logits (cho Human và AI).
    
    # Nếu output là logits cho 2 lớp (ví dụ: [logit_human, logit_ai])
    logits = outputs.logits
    probabilities = F.softmax(logits, dim=1) # Áp dụng softmax để có xác suất
    
    # Lấy xác suất của lớp "AI" (giả sử lớp 1 là AI)
    # Bạn cần xác nhận lại id của lớp AI từ quá trình huấn luyện của bạn
    # Ví dụ: nếu model.config.id2label = {0: 'human', 1: 'ai'}
    ai_probability = probabilities[:, 1].item() * 100 # Xác suất của lớp AI
    human_probability = probabilities[:, 0].item() * 100 # Xác suất của lớp Human

    return f"AI Percentage: {ai_probability:.2f}% (Human Percentage: {human_probability:.2f}%)"


# 3. Tạo giao diện Gradio
# gr.Interface(fn, inputs, outputs, title, description)
# fn: Hàm dự đoán
# inputs: Các thành phần đầu vào (ví dụ: gr.Textbox)
# outputs: Các thành phần đầu ra (ví dụ: gr.Textbox)
# title, description: Tiêu đề và mô tả cho giao diện

iface = gr.Interface(
    fn=predict_ai_human,
    inputs=gr.Textbox(lines=10, label="Nhập văn bản vào đây để kiểm tra"),
    outputs=gr.Textbox(label="Kết quả dự đoán"),
    title="Phân loại văn bản AI/Human",
    description="Một demo nhỏ để kiểm tra tỷ lệ văn bản được tạo bởi AI so với văn bản của con người."
)

# 4. Khởi chạy giao diện
if __name__ == "__main__":
    iface.launch()