infer working locally - first try at HF

app.py

@@ -0,0 +1,54 @@
+import os
+import time
+import gradio as gr
+import numpy as np
+from huggingface_hub import hf_hub_download
+
+from src.infer import load_model, predict
+
+os.environ.setdefault("HF_HOME", "/data/.huggingface")
+
+_model = None
+_ckpt_path = None
+
+def _warmup():
+    global _model, _ckpt_path
+    if _model is not None:
+        return
+    t0 = time.time()
+    _ckpt_path = hf_hub_download(
+        repo_id="rhasan/empathy",
+        filename="UPLME_NewsEmp_tuned-lambdas.ckpt",
+        repo_type="model",
+        local_dir="/data/uplme_ckpt"
+    )
+    load_model(_ckpt_path)
+    return f"Model loaded in {time.time() - t0:.1f} seconds."
+
+def predict_with_ci(essay: str, article: str) -> dict:
+    _warmup()
+    mean, var = predict(essay, article)
+    # scores were originally in [1, 7]
+    # lets scale them to [0, 100]
+    mean = (mean - 1) / 6 * 100
+
+    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)}
+
+with gr.Blocks(title="Empathy Prediction") 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...")
+            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])
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,6 @@
+torch
+transformers
+gradio
+lightning
+numpy
+huggingface_hub

src/infer.py

@@ -0,0 +1,38 @@
+"""
+FROM https://github.com/hasan-rakibul/UPLME/tree/main
+"""
+
+import torch
+from transformers import AutoTokenizer
+
+from paired_texts_modelling import LitPairedTextModel
+
+_device = None
+_model = None
+_tokeniser = None
+
+def load_model(ckpt_path: str):
+    global _model, _tokeniser, _device
+    plm_name = "roberta-base"
+    _device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    _model = LitPairedTextModel.load_from_checkpoint(ckpt_path).to(_device).eval()
+    _tokeniser = AutoTokenizer.from_pretrained(
+        plm_name,
+        use_fast=True,
+        add_prefix_space=False # the first word is tokenised differently if not a prefix space, but it might decrease performance, so False (09/24)
+    )
+
+@torch.inference_mode()
+def predict(essay: str, article: str) -> tuple[float, float]:
+    max_length = 512
+    toks = _tokeniser(
+        essay,
+        article,
+        truncation=True,
+        max_length=max_length,
+        return_tensors="pt"
+    ).to(_device)
+    mean, var, _ = _model(toks)
+    return mean.item(), var.item()
+
+    

src/paired_texts_modelling.py

@@ -0,0 +1,131 @@
+"""
+FROM https://github.com/hasan-rakibul/UPLME/tree/main
+"""
+
+import torch
+from torch import Tensor
+import lightning as L
+from transformers import (
+    AutoModel, 
+)
+import logging
+
+import lightning as L
+
+
+logger = logging.getLogger(__name__)
+    
+class CrossEncoderProbModel(torch.nn.Module):
+    def __init__(self, plm_name: str):
+        super().__init__()
+        self.model = AutoModel.from_pretrained(plm_name)
+
+        if plm_name.startswith("roberta"):
+            # only applicable for roberta
+            self.pooling = "roberta-pooler"
+        else:
+            self.pooling = "cls"
+
+        self.out_proj_m = torch.nn.Sequential(
+            torch.nn.LayerNorm(self.model.config.hidden_size),
+            torch.nn.Dropout(0.25),
+            torch.nn.Linear(self.model.config.hidden_size, 1)
+        )
+
+        self.out_proj_v = torch.nn.Sequential(
+            torch.nn.LayerNorm(self.model.config.hidden_size),
+            torch.nn.Dropout(0.25),
+            torch.nn.Linear(self.model.config.hidden_size, 1),
+            torch.nn.Softplus()
+        )
+
+    def forward(self, input_ids, attention_mask):
+        output = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask
+        )
+
+        if self.pooling == "mean":
+            sentence_representation = (
+                (output.last_hidden_state * attention_mask.unsqueeze(-1)).sum(-2) /
+                attention_mask.sum(dim=-1).unsqueeze(-1)
+            )
+        elif self.pooling == "cls":
+            sentence_representation = output.last_hidden_state[:, 0, :]
+        elif self.pooling == "roberta-pooler":
+            sentence_representation = output.pooler_output # (batch_size, hidden_dim)
+
+        mean = self.out_proj_m(sentence_representation)
+        var = self.out_proj_v(sentence_representation)
+        var = torch.clamp(var, min=1e-8, max=1000) # following Seitzer-NeurIPS2022
+
+        return mean.squeeze(), var.squeeze(), sentence_representation, output.last_hidden_state
+    
+class LitPairedTextModel(L.LightningModule):
+    def __init__(
+        self,
+        plm_names: list[str],
+        lr: float,
+        log_dir: str,
+        save_uc_metrics: bool,
+        error_decay_factor: float,
+        approach: str,
+        sep_token_id: int, # required for alignment loss
+        lambdas: list[float] = [], # initlisaed to compatible with old saved checkpoints
+        num_passes: int = 4
+    ):
+        super().__init__()
+        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.lr = lr
+        self.log_dir = log_dir
+        self.save_uc_metrics = save_uc_metrics
+
+        self.error_decay_factor = error_decay_factor
+        
+        self.lambdas = lambdas
+        self.sep_token_id = sep_token_id
+        self.num_passes = num_passes
+        
+        self.penalty_type = "exp-decay"
+
+        self.validation_outputs = []
+        self.test_outputs = []
+    
+    def forward(self, batch: dict) -> tuple[Tensor, Tensor, Tensor]:
+        means, varss, hidden_states = [], [], []
+
+        for _ in range(self.num_passes):
+            if self.approach == "cross-prob":
+                mean, var, _, hidden_state = self.model(
+                    input_ids=batch['input_ids'],
+                    attention_mask=batch['attention_mask']
+                )
+            elif self.approach == "cross-basic":
+                mean, hidden_state = self.model(batch)
+                var = torch.zeros_like(mean)
+            
+            means.append(mean)
+            varss.append(var)
+            hidden_states.append(hidden_state)
+        
+        mean = torch.stack(means, dim=0).mean(dim=0)
+        var = torch.stack(varss, dim=0).mean(dim=0)
+        hidden_state = torch.stack(hidden_states, dim=0).mean(dim=0)
+
+        return mean, var, hidden_state
+    
+    def _enable_dropout_at_inference(self):
+        for m in self.model.modules():
+            if isinstance(m, torch.nn.Dropout):
+                m.train()
+
+