# -*- coding: utf-8 -*-
"""app.py

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1eyNEXhQE4T_7cq-MsPQ77p7h6xdrOpzk
"""

import gradio as gr
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# 모델 경로 설정
model_path = "./model"  # 업로드된 모델 디렉토리 경로

# 모델과 토크나이저 로드
model = AutoModelForSequenceClassification.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained("klue/bert-base")

# 예측 함수
def predict(text):
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model(**inputs)
    probabilities = torch.sigmoid(outputs.logits)
    depression_prob = probabilities[0, 1].item()

    if depression_prob > 0.5:
        return f"Depressed (Confidence: {depression_prob:.2%})"
    else:
        return f"Not Depressed (Confidence: {1 - depression_prob:.2%})"

# Gradio 인터페이스
interface = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(label="Enter your text here"),
    outputs=gr.Textbox(label="Result"),
    title="Depression Detection",
    description="Predict the likelihood of depression based on text input.",
)

interface.launch()