add real code

blaser_2_0_qe.py

@@ -15,6 +15,10 @@
 
 import evaluate
 import datasets
+import torch
+
+from sonar.inference_pipelines.text import TextToEmbeddingModelPipeline
+from sonar.models.blaser.loader import load_blaser_model
 
 
 # TODO: Add BibTeX citation
@@ -71,8 +75,8 @@ class BLASER20QE(evaluate.Metric):
             inputs_description=_KWARGS_DESCRIPTION,
             # This defines the format of each prediction and reference
             features=datasets.Features({
-                'predictions': datasets.Value('int64'),
-                'references': datasets.Value('int64'),
+                'predictions': datasets.Value('str'),
+                'references': datasets.Value('str'),
             }),
             # Homepage of the module for documentation
             homepage="http://module.homepage",
@@ -83,13 +87,21 @@ class BLASER20QE(evaluate.Metric):
 
     def _download_and_prepare(self, dl_manager):
         """Optional: download external resources useful to compute the scores"""
-        # TODO: Download external resources if needed
-        pass
+        
+        self.text_embedder = TextToEmbeddingModelPipeline(encoder="text_sonar_basic_encoder", tokenizer="text_sonar_basic_encoder")
+        self.blaser_qe = load_blaser_model("blaser_2_0_qe").eval()
+        #self.blaser_ref = load_blaser_model("blaser_2_0_ref").eval()
 
-    def _compute(self, predictions, references):
+    def _compute(self, sources, translations):
         """Returns the scores"""
-        # TODO: Compute the different scores of the module
-        accuracy = sum(i == j for i, j in zip(predictions, references)) / len(predictions)
+        # TODO: adjust the languages
+        src_embs = text_embedder.predict(sources, source_lang="eng_Latn")
+        mt_embs = text_embedder.predict(translations, source_lang="eng_Latn")
+
+        with torch.inference_mode():
+            #ref_score = blaser_ref(src=src_embs, ref=ref_embs, mt=mt_embs).mean().item()
+            qe_score = blaser_qe(src=src_embs, mt=mt_embs).mean().item()
+
         return {
-            "accuracy": accuracy,
+            "BLASER 2.0-QE": qe_score,
         }

requirements.txt

@@ -1 +1,2 @@
-git+https://github.com/huggingface/evaluate@main
+git+https://github.com/huggingface/evaluate@main
+sonar-space