app.py

import gradio as gr
import torch
from transformers import BertTokenizer, BertForSequenceClassification, XLNetTokenizer, XLNetForSequenceClassification

# โหลดโมเดลของคุณ
logbert_model = BertForSequenceClassification.from_pretrained("Sirapatsorn/Spark_Log_Analysis", use_auth_token=True)
xlnet_model = XLNetForSequenceClassification.from_pretrained("Sirapatsorn/Spark_Log_Analysis", use_auth_token=True)
logbert_tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
xlnet_tokenizer = XLNetTokenizer.from_pretrained("xlnet-base-cased")

# ตรวจสอบว่ามี GPU หรือไม่
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logbert_model.to(device)
xlnet_model.to(device)

# ฟังก์ชันการพยากรณ์
def predict_log(text):
    logbert_inputs = logbert_tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(device)
    with torch.no_grad():
        logbert_outputs = logbert_model(**logbert_inputs)
        log_level = torch.argmax(logbert_outputs.logits, dim=1).item()
        log_level_confidence = torch.softmax(logbert_outputs.logits, dim=1)[0][log_level].item()

    log_levels = ["INFO", "WARN", "ERROR"]
    log_level_result = log_levels[log_level]

    xlnet_inputs = xlnet_tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(device)
    with torch.no_grad():
        xlnet_outputs = xlnet_model(**xlnet_inputs)
        performance_value = xlnet_outputs.logits.item()

    if performance_value < 0:
        performance_status = "Good Performance"
    elif performance_value < 3.0:
        performance_status = "Normal Performance"
    else:
        performance_status = "Poor Performance"

    return {
        "Log Level": log_level_result,
        "Confidence": f"{log_level_confidence:.2f}",
        "Performance Value": f"{performance_value:.2f}",
        "Performance Status": performance_status
    }

# ฟังก์ชันสำหรับ Gradio: พยากรณ์จากข้อความโดยตรง
def predict_from_text(text):
    prediction = predict_log(text)
    return (f"Log Level: {prediction['Log Level']} (Confidence: {prediction['Confidence']})\n"
            f"Performance Value: {prediction['Performance Value']}\n"
            f"Performance Status: {prediction['Performance Status']}")

def predict_from_file(file):
    results = []
    with open(file.name, 'r') as f:
        for line in f:
            prediction = predict_log(line.strip())
            result_text = (f"Log: {line.strip()}\n"
                           f"Log Level: {prediction['Log Level']} (Confidence: {prediction['Confidence']})\n"
                           f"Performance Value: {prediction['Performance Value']}\n"
                           f"Performance Status: {prediction['Performance Status']}")
            results.append(result_text)
    return "\n\n".join(results)

custom_css = """
.gr-button {
    background-color: #FFA500 !important; 
    color: #FFFFFF !important;
    border: none !important;
}
"""

with gr.Blocks(css=custom_css) as demo:
    with gr.Tabs():
        with gr.TabItem("Upload File"):
            file_upload = gr.File(label="Upload Log File")
            file_btn = gr.Button("Predict")
            file_output = gr.Textbox(label="Output")
            file_btn.click(predict_from_file, inputs=file_upload, outputs=file_output)
        with gr.TabItem("Text Input"):
            text_input = gr.Textbox(label="Enter Log Message")
            text_btn = gr.Button("Predict")
            text_output = gr.Textbox(label="Output")
            text_btn.click(predict_from_text, inputs=text_input, outputs=text_output)

    demo.css += """
    .gr-button.gr-button-lg.gr-button-secondary {
        display: none !important;
    }
    """

demo.launch()