fixed module load

app.py

@@ -17,7 +17,7 @@ def _warmup():
         return
     t0 = time.time()
     _ckpt_path = hf_hub_download(
-        repo_id="rhasan/empathy",
+        repo_id="rhasan/UPLME",
         filename="UPLME_NewsEmp_tuned-lambdas.ckpt",
         repo_type="model",
         local_dir="/data/uplme_ckpt"
@@ -25,7 +25,7 @@ def _warmup():
     load_model(_ckpt_path)
     return f"Model loaded in {time.time() - t0:.1f} seconds."
 
-def predict_with_ci(essay: str, article: str) -> dict:
+def predict_with_ci(essay: str, article: str) -> tuple[float, float, float]:
     _warmup()
     mean, var = predict(essay, article)
     # scores were originally in [1, 7]
@@ -35,7 +35,7 @@ def predict_with_ci(essay: str, article: str) -> dict:
     std = np.sqrt(var)
     ci_low = max(0.0, mean - 1.96 * std)
     ci_upp = min(100.0, mean + 1.96 * std)
-    return {"mean": mean, "ci": (ci_low, ci_upp)}
+    return mean, ci_low, ci_upp
 
 with gr.Blocks(title="Empathy Prediction") as demo:
     gr.Markdown("# Empathy Prediction with Uncertainty Estimation")
@@ -45,10 +45,11 @@ with gr.Blocks(title="Empathy Prediction") as demo:
             article_input = gr.Textbox(label="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=4)
-            ci = gr.Number(label="95\% CI", precision=4)
-        
-        button.click(fn=predict_with_ci, inputs=[essay_input, article_input], outputs=[output_mean, ci])
+            output_mean = gr.Number(label="Predicted Empathy Mean", precision=2)
+            ci_low = gr.Number(label="95% CI Lower Bound", precision=2)
+            ci_upp = gr.Number(label="95% CI Upper Bound", precision=2)
+
+        button.click(fn=predict_with_ci, inputs=[essay_input, article_input], outputs=[output_mean, ci_low, ci_upp])
 
 if __name__ == "__main__":
     demo.launch()

src/infer.py

@@ -5,7 +5,7 @@ FROM https://github.com/hasan-rakibul/UPLME/tree/main
 import torch
 from transformers import AutoTokenizer
 
-from paired_texts_modelling import LitPairedTextModel
+from src.paired_texts_modelling import LitPairedTextModel
 
 _device = None
 _model = None

src/paired_texts_modelling.py

@@ -10,9 +10,6 @@ from transformers import (
 )
 import logging
 
-import lightning as L
-
-
 logger = logging.getLogger(__name__)
     
 class CrossEncoderProbModel(torch.nn.Module):
@@ -78,12 +75,7 @@ class LitPairedTextModel(L.LightningModule):
         self.save_hyperparameters()
 
         self.approach = approach
-        if self.approach == "cross-basic":
-            self.model = CrossEncoderBasicModel(plm_name=plm_names[0])
-        elif self.approach == "cross-prob":
-            self.model = CrossEncoderProbModel(plm_name=plm_names[0])
-        else:
-            raise ValueError(f"Invalid approach: {self.approach}")
+        self.model = CrossEncoderProbModel(plm_name=plm_names[0])
 
         self.lr = lr
         self.log_dir = log_dir
@@ -101,6 +93,7 @@ class LitPairedTextModel(L.LightningModule):
         self.test_outputs = []
     
     def forward(self, batch: dict) -> tuple[Tensor, Tensor, Tensor]:
+        self._enable_dropout_at_inference()
         means, varss, hidden_states = [], [], []
 
         for _ in range(self.num_passes):