fix roberta on sescore.py

__init__.py

@@ -1,38 +1 @@
-import comet
-from typing import Dict
-import torch
-from comet.encoders.base import Encoder
-from comet.encoders.bert import BERTEncoder
-from transformers import AutoModel, AutoTokenizer
 
-class robertaEncoder(BERTEncoder):
-    def __init__(self, pretrained_model: str) -> None:
-        super(Encoder, self).__init__()
-        self.tokenizer = AutoTokenizer.from_pretrained(pretrained_model)
-        self.model = AutoModel.from_pretrained(
-            pretrained_model, add_pooling_layer=False
-        )
-        self.model.encoder.output_hidden_states = True
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model: str) -> Encoder:
-        return robertaEncoder(pretrained_model)
-
-    def forward(
-        self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs
-    ) -> Dict[str, torch.Tensor]:
-        last_hidden_states, _, all_layers = self.model(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            output_hidden_states=True,
-            return_dict=False,
-        )
-        return {
-            "sentemb": last_hidden_states[:, 0, :],
-            "wordemb": last_hidden_states,
-            "all_layers": all_layers,
-            "attention_mask": attention_mask,
-        }
-
-# initialize roberta into str2encoder
-comet.encoders.str2encoder['RoBERTa'] = robertaEncoder

sescore.py

@@ -16,6 +16,42 @@
 import evaluate
 import datasets
 
+import comet
+from typing import Dict
+import torch
+from comet.encoders.base import Encoder
+from comet.encoders.bert import BERTEncoder
+from transformers import AutoModel, AutoTokenizer
+
+class robertaEncoder(BERTEncoder):
+    def __init__(self, pretrained_model: str) -> None:
+        super(Encoder, self).__init__()
+        self.tokenizer = AutoTokenizer.from_pretrained(pretrained_model)
+        self.model = AutoModel.from_pretrained(
+            pretrained_model, add_pooling_layer=False
+        )
+        self.model.encoder.output_hidden_states = True
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model: str) -> Encoder:
+        return robertaEncoder(pretrained_model)
+
+    def forward(
+        self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs
+    ) -> Dict[str, torch.Tensor]:
+        last_hidden_states, _, all_layers = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True,
+            return_dict=False,
+        )
+        return {
+            "sentemb": last_hidden_states[:, 0, :],
+            "wordemb": last_hidden_states,
+            "all_layers": all_layers,
+            "attention_mask": attention_mask,
+        }
+
 
 # TODO: Add BibTeX citation
 _CITATION = """\
@@ -87,12 +123,14 @@ class SEScore(evaluate.Metric):
         from comet import load_from_checkpoint
         import gdown
         import os
+        # initialize roberta into str2encoder
+        comet.encoders.str2encoder['RoBERTa'] = robertaEncoder
         url = "https://drive.google.com/uc?id=1QgMP_Y4QCbvDMTeVacYt0J76OYvwWK9V&export=download&confirm=true"
         output = 'sescore_ckpt.gz'
         gdown.download(url, output, quiet=False)
         cmd = 'tar -xvf sescore_ckpt.gz'
         os.system(cmd)
-        self.scorer = load_from_checkpoint('sescore_ckpt/zh_en/checkpoint/sescore_english.ckpt')
+        self.scorer = load_from_checkpoint('/home/user/app/sescore_ckpt/zh_en/checkpoint/sescore_english.ckpt')
 
     def _compute(self, sources, predictions, references, gpus=None, progress_bar=False):
         if gpus is None: