Update app.py

app.py

@@ -3,33 +3,42 @@ import torch
 import os
 from transformers import BertTokenizer, BertForSequenceClassification, XLNetTokenizer, XLNetForSequenceClassification
 
-# โหลดโมเดลของคุณ
-logbert_model = BertForSequenceClassification.from_pretrained("Sirapatsorn/Spark_Log_Analysis", token=os.getenv("HUGGINGFACE_TOKEN"))
-xlnet_model = XLNetForSequenceClassification.from_pretrained("Sirapatsorn/Spark_Log_Analysis", token=os.getenv("HUGGINGFACE_TOKEN"))
+# Load models
+logbert_model = BertForSequenceClassification.from_pretrained(
+    "Sirapatsorn/Spark_Log_Analysis-logbert",  # แก้เป็น path ที่ถูกต้อง
+    token=os.getenv("HUGGINGFACE_TOKEN")
+)
+xlnet_model = XLNetForSequenceClassification.from_pretrained(
+    "Sirapatsorn/Spark_Log_Analysis-xlnet",    # แก้เป็น path ที่ถูกต้อง
+    token=os.getenv("HUGGINGFACE_TOKEN")
+)
+
+# Load tokenizers
 logbert_tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
 xlnet_tokenizer = XLNetTokenizer.from_pretrained("xlnet-base-cased")
 
-# ตรวจสอบว่ามี GPU หรือไม่
+# Check for 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 prediction
     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_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 prediction
     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()
-
+    performance_value = xlnet_outputs.logits.item()
+    
+    # Determine performance status
     if performance_value < 0:
         performance_status = "Good Performance"
     elif performance_value < 3.0:
@@ -44,7 +53,6 @@ def predict_log(text):
         "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"
@@ -57,31 +65,36 @@ def predict_from_file(file):
         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']}")
+                         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)
 
+# Create Gradio interface
 custom_css = """
 .gr-button {
-    background-color: #FFA500 !important; 
+    background-color: #FFA500 !important;
     color: #FFFFFF !important;
     border: none !important;
 }
 """
 
 with gr.Blocks(css=custom_css) as demo:
+    gr.Markdown("# Spark Log Analysis")
+    gr.Markdown("Analyze your log messages for log level and performance prediction")
+    
     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_output = gr.Textbox(label="Output", lines=10)
             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_input = gr.Textbox(label="Enter Log Message", lines=3)
             text_btn = gr.Button("Predict")
-            text_output = gr.Textbox(label="Output")
+            text_output = gr.Textbox(label="Output", lines=5)
             text_btn.click(predict_from_text, inputs=text_input, outputs=text_output)
 
     demo.css += """
@@ -90,4 +103,5 @@ with gr.Blocks(css=custom_css) as demo:
     }
     """
 
-demo.launch()
+# Launch the app
+demo.launch()