adding fig

app.py

@@ -1,31 +1,43 @@
 import os
+from pathlib import Path
 import time
 import gradio as gr
+from gadio.themes import Soft
 import numpy as np
+import matplotlib.pyplot as plt
 from huggingface_hub import hf_hub_download
 
 from src.infer import load_model, predict
 
-os.environ.setdefault("HF_HOME", "/data/.huggingface")
+os.environ.setdefault("HF_HOME", str(Path.home() / ".cache" / "huggingface"))
 
 _model = None
-_ckpt_path = None
 
 def _warmup():
-    global _model, _ckpt_path
+    global _model
     if _model is not None:
         return
     t0 = time.time()
-    _ckpt_path = hf_hub_download(
+    ckpt_path = hf_hub_download(
         repo_id="rhasan/UPLME",
         filename="UPLME_NewsEmp_tuned-lambdas.ckpt",
-        repo_type="model",
-        local_dir="/data/uplme_ckpt"
+        repo_type="model"
     )
-    load_model(_ckpt_path)
+    load_model(ckpt_path)
     return f"Model loaded in {time.time() - t0:.1f} seconds."
 
-def predict_with_ci(essay: str, article: str) -> tuple[float, float, float]:
+def ci_plot(mean: float, low: float, upp: float):
+    fig, ax = plt.subplots(figsize=(6, 1))
+    ax.hlines(1, 0, 100, linewidht=2, alpha=0.15)
+    ax.hlines(1, low, upp, linewidht=6)
+    ax.plot([mean], [1], "o")
+    ax.set_xlim(0, 100)
+    ax.set_yticks([])
+    ax.set_xlabel("Empathy Score (0-100)")
+    fig.tight_layout()
+    return fig
+
+def predict_with_ci(essay: str, article: str) -> tuple[float, float, float, plt.Figure]:
     _warmup()
     mean, var = predict(essay, article)
     # scores were originally in [1, 7]
@@ -35,21 +47,34 @@ def predict_with_ci(essay: str, article: str) -> tuple[float, float, float]:
     std = np.sqrt(var)
     ci_low = max(0.0, mean - 1.96 * std)
     ci_upp = min(100.0, mean + 1.96 * std)
-    return mean, ci_low, ci_upp
+    fig = ci_plot(mean, ci_low, ci_upp)
+    return mean, ci_low, ci_upp, fig
+
 
-with gr.Blocks(title="Empathy Prediction") as demo:
+with gr.Blocks(title="UPLME", theme=Soft(primary_hue="blue")) as demo:
     gr.Markdown("# Empathy Prediction with Uncertainty Estimation")
     with gr.Row():
         with gr.Column():
-            essay_input = gr.Textbox(label="Essay", lines=10, placeholder="Enter the essay text here...")
-            article_input = gr.Textbox(label="Article", lines=10, placeholder="Enter the article text here...")
+            essay_input = gr.Textbox(label="Response (E.g., Essay)", lines=10, placeholder="Enter the essay text here...")
+            article_input = gr.Textbox(label="Stimulus (E.g., News Article)", lines=10, placeholder="Enter the article text here...")
             button = gr.Button("Predict")
         with gr.Column():
-            output_mean = gr.Number(label="Predicted Empathy Mean", precision=2)
+            output_mean = gr.Number(label="Predicted Empathy Score", precision=2)
             ci_low = gr.Number(label="95% CI Lower Bound", precision=2)
             ci_upp = gr.Number(label="95% CI Upper Bound", precision=2)
+        
+        fig = gr.Plot(label="Prediction +/- 95% CI")
+
+        button.click(fn=predict_with_ci, inputs=[essay_input, article_input], outputs=[output_mean, ci_low, ci_upp, fig])
 
-        button.click(fn=predict_with_ci, inputs=[essay_input, article_input], outputs=[output_mean, ci_low, ci_upp])
+    gr.Markdown("## About")
+    gr.Markdown("""
+    This application predicts empathy score using the UPLME model proposed in **UPLME: Uncertainty-Aware Probabilistic Language Modelling for Robust Empathy Regression** by **Md Rakibul Hasan, Md Zakir Hossain, Aneesh Krishna, Shafin Rahman and Tom Gedeon**.
+    The model provides both a mean empathy score and uncertainty estimates.
+    - **Response**: The text input representing the individual's response (e.g., essay).
+    - **Stimulus**: The text input representing the stimulus (e.g. newspaper article) that the response is based on.
+    - **Predicted Empathy**: The predicted empathy score on a scale from 0 to 100.
+    """)
 
 if __name__ == "__main__":
     demo.launch()

src/paired_texts_modelling.py

@@ -120,5 +120,3 @@ class LitPairedTextModel(L.LightningModule):
         for m in self.model.modules():
             if isinstance(m, torch.nn.Dropout):
                 m.train()
-
-