Update README.md

README.md

@@ -13,67 +13,126 @@ license: cc-by-4.0
 pipeline_tag: translation
 ---
 
-# NorT5 large finetuned for English -> Norwegian Bokmål translation
+# NorT5 large finetuned for English → Norwegian (Bokmål or Nynorsk) translation
 
 <img src="https://huggingface.co/ltg/norbert3-base/resolve/main/norbert.png" width=12.5%>
 
-The official release of a new generation of NorT5 language models described in paper [**NorBench — A Benchmark for Norwegian Language Models**](https://arxiv.org/abs/2305.03880). Plese read the paper to learn more details about the model.
-
-
-## Other sizes:
-- [NorT5 xs (32M)](https://huggingface.co/ltg/nort5-xs)
-- [NorT5 small (88M)](https://huggingface.co/ltg/nort5-small)
-- [NorT5 base (228M)](https://huggingface.co/ltg/nort5-base)
-- [NorT5 large (808M)](https://huggingface.co/ltg/nort5-large)
-
-
-## Encoder-only NorBERT siblings:
-- [NorBERT 3 xs (15M)](https://huggingface.co/ltg/norbert3-xs)
-- [NorBERT 3 small (40M)](https://huggingface.co/ltg/norbert3-small)
-- [NorBERT 3 base (123M)](https://huggingface.co/ltg/norbert3-base)
-- [NorBERT 3 large (323M)](https://huggingface.co/ltg/norbert3-large)
 
 
 ## Example usage
 
-This model currently needs a custom wrapper from `modeling_nort5.py`, you should therefore load the model with `trust_remote_code=True`.
+This model is specifically finetuned for translating documents in the English-to-Norwegian direction.
+Unlike traditional NMT models, it is trained on paragraph-to-paragraph translation – the translation quality is thus better if you feed it whole paragraphs instead of segmented sentences.
 
+A simple example of how to use this model can be found in the `translate.py` file:
 
 ```python
 import torch
+import transformers
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers.generation import LogitsProcessor
+
+
+class RepetitionPenaltyLogitsProcessor(LogitsProcessor):
+    def __init__(self, penalty: float, model):
+        last_bias = model.classifier.nonlinearity[-1].bias.data
+        last_bias = torch.nn.functional.log_softmax(last_bias)
+        self.penalty = penalty * (last_bias - last_bias.max())
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        penalized_score = torch.gather(scores + self.penalty.unsqueeze(0).to(input_ids.device), 1, input_ids).to(scores.dtype)
+        scores.scatter_(1, input_ids, penalized_score)
+        return scores
+
+
+class Translator:
+    def __init__(self, model_path="ltg/nort5-large-en-no-translation", device="cpu"):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.cls_index = self.tokenizer.convert_tokens_to_ids("[CLS]")
+        self.sep_index = self.tokenizer.convert_tokens_to_ids("[SEP]")
+        self.eos_index = self.tokenizer.convert_tokens_to_ids("[EOS]")
+        self.pad_index = self.tokenizer.convert_tokens_to_ids("[PAD]")
+        self.eng_index = self.tokenizer.convert_tokens_to_ids(">>eng<<")
+        self.nob_index = self.tokenizer.convert_tokens_to_ids(">>nob<<")
+        self.nno_index = self.tokenizer.convert_tokens_to_ids(">>nno<<")
+
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_path, trust_remote_code=True)
+
+        self.device = device
+        print(f"SYSTEM: Running on {self.device}", flush=True)
+
+        self.model = self.model.to(device)
+        self.model.eval()
+
+        print(f"Sucessfully loaded the model to the memory")
+
+        self.LANGUAGE_IDS = {
+            "en": self.eng_index,
+            "nb": self.nob_index,
+            "nn": self.nno_index
+        }
+
+    def __call__(self, source, source_language, target_language):
+        source = [s.strip() for s in source.split('\n')]
+        source_subwords = self.tokenizer(source).input_ids
+        source_subwords = [[self.cls_index, self.LANGUAGE_IDS[target_language], self.LANGUAGE_IDS[source_language]] + s + [self.sep_index] for s in source_subwords]
+        source_subwords = [torch.tensor(s) for s in source_subwords]
+        source_subwords = torch.nn.utils.rnn.pad_sequence(source_subwords, batch_first=True, padding_value=self.pad_index)
+        source_subwords = source_subwords[:, :512].to(self.device)
+
+        def generate(model, **kwargs):
+            with torch.inference_mode():
+                with torch.autocast(enabled=self.device != "cpu", device_type="cuda", dtype=torch.bfloat16):
+                    return model.generate(**kwargs)
+
+        generate_kwargs = dict(
+            input_ids=source_subwords,
+            attention_mask=(source_subwords != self.pad_index).long(),
+            max_new_tokens = 512-1,
+            num_beams=8,
+            length_penalty=1.6,
+            early_stopping=True,
+            do_sample=False,
+            use_cache=True,
+            logits_processor=[RepetitionPenaltyLogitsProcessor(0.5, self.model), transformers.LogitNormalization()]
+        )
+        output = generate(self.model, **generate_kwargs).tolist()
+        paragraphs = [self.tokenizer.decode(c, skip_special_tokens=True).strip() for c in output]
+        translation = '\n'.join(paragraphs)
+
+        return translation
+
+
+if __name__ == "__main__":
+
+    translator = Translator()
+
+    en_text = "How are you feeling right now? Better?"
+    no_text = translator(en_text, "en", "nb")
+
+    print(en_text)
+    print(no_text)
+```
 
-tokenizer = AutoTokenizer.from_pretrained("ltg/nort5-large")
-model = AutoModelForSeq2SeqLM.from_pretrained("ltg/nort5-large", trust_remote_code=True)
-
-
-# MASKED LANGUAGE MODELING
-
-sentence = "Brukseksempel: Elektrisk oppvarming. Definisjonen på ordet oppvarming er[MASK_0]."
-encoding = tokenizer(sentence)
-
-input_tensor = torch.tensor([encoding.input_ids])
-output_tensor = model.generate(input_tensor, decoder_start_token_id=7, eos_token_id=8)
-tokenizer.decode(output_tensor.squeeze(), skip_special_tokens=True)
-
-# should output:  å varme opp
 
+## The NorT5 and NorBERT family
 
-# PREFIX LANGUAGE MODELING
-# you need to finetune this model or use `nort5-{size}-lm` model, which is finetuned on prefix language modeling
+The official release of a new generation of NorT5 language models described in paper [**NorBench — A Benchmark for Norwegian Language Models**](https://arxiv.org/abs/2305.03880). Plese read the paper to learn more details about the model.
 
-sentence = "Brukseksempel: Elektrisk oppvarming. Definisjonen på ordet oppvarming er (Wikipedia) "
-encoding = tokenizer(sentence)
 
-input_tensor = torch.tensor([encoding.input_ids])
-output_tensor = model.generate(input_tensor, max_new_tokens=50, num_beams=4, do_sample=False)
-tokenizer.decode(output_tensor.squeeze())
+## Other sizes:
+- [NorT5 xs (32M)](https://huggingface.co/ltg/nort5-xs)
+- [NorT5 small (88M)](https://huggingface.co/ltg/nort5-small)
+- [NorT5 base (228M)](https://huggingface.co/ltg/nort5-base)
+- [NorT5 large (808M)](https://huggingface.co/ltg/nort5-large)
 
-# should output: [BOS]ˈoppvarming, det vil si at det skjer en endring i temperaturen i et medium, f.eks. en ovn eller en radiator, slik at den blir varmere eller kaldere, eller at den blir varmere eller kaldere, eller at den blir
-```
 
+## Encoder-only NorBERT siblings:
+- [NorBERT 3 xs (15M)](https://huggingface.co/ltg/norbert3-xs)
+- [NorBERT 3 small (40M)](https://huggingface.co/ltg/norbert3-small)
+- [NorBERT 3 base (123M)](https://huggingface.co/ltg/norbert3-base)
+- [NorBERT 3 large (323M)](https://huggingface.co/ltg/norbert3-large)
 
-The following classes are currently implemented: `AutoModel`, `AutoModelForSeq2SeqLM`.
 
 ## Cite us