isort/black

README.md

@@ -67,8 +67,9 @@ After importing and loading the model and a prefix tree (trie), you would genera
 
 ```python
 import pickle
-from genre.trie import Trie
+
 from genre.fairseq_model import GENRE
+from genre.trie import Trie
 
 # load the prefix tree (trie)
 with open("../data/kilt_titles_trie_dict.pkl", "rb") as f:
@@ -97,8 +98,9 @@ Making predictions with mGENRE is very similar, but we additionally need to map
 
 ```python
 import pickle
-from genre.trie import Trie, MarisaTrie
+
 from genre.fairseq_model import mGENRE
+from genre.trie import MarisaTrie, Trie
 
 with open("../data/lang_title2wikidataID-normalized_with_redirect.pkl", "rb") as f:
     lang_title2wikidataID = pickle.load(f)

genre/fairseq_model.py

@@ -104,7 +104,7 @@ def encode(self, sentence) -> torch.LongTensor:
         else:
             return tokens
 
-    
+
 class GENRE(BARTModel):
     @classmethod
     def from_pretrained(

genre/hf_model.py

@@ -8,9 +8,10 @@
 from typing import Dict, List
 
 import torch
-from genre.utils import chunk_it
 from transformers import BartForConditionalGeneration, BartTokenizer
 
+from genre.utils import chunk_it
+
 logger = logging.getLogger(__name__)
 
 

requirements-test.txt

@@ -1,3 +1,3 @@
 torch
 requests
-pytest
+pytest

scripts_genre/convert_kilt_to_fairseq.py

@@ -26,7 +26,14 @@ def convert_kilt_to_fairseq(dataset):
             for prov in out["provenance"]
             if prov.get("bleu_score", 1) > 0.5
         ):
-            source.append(create_input(doc, max_length=384, start_delimiter="[START_ENT]", end_delimiter="[END_ENT]"))
+            source.append(
+                create_input(
+                    doc,
+                    max_length=384,
+                    start_delimiter="[START_ENT]",
+                    end_delimiter="[END_ENT]",
+                )
+            )
             target.append(title)
             if "meta" in doc and "template_questions" in doc["meta"]:
                 for template_question in doc["meta"]["template_questions"]:

scripts_mgenre/preprocess_sentencepiece.py

@@ -23,12 +23,12 @@ def initializer(self):
         sp = spm.SentencePieceProcessor(model_file=self.args.model)
         old2new = None
         if self.args.product_vocab_size is not None:
-            assert sp.vocab_size() <= self.args.product_vocab_size ** 2
+            assert sp.vocab_size() <= self.args.product_vocab_size**2
             rand = random.Random(self.args.seed)
             old2new = [
                 (x // self.args.product_vocab_size, x % self.args.product_vocab_size)
                 for x in rand.sample(
-                    range(self.args.product_vocab_size ** 2), sp.vocab_size()
+                    range(self.args.product_vocab_size**2), sp.vocab_size()
                 )
             ]
 

tests/test_example.py

@@ -3,8 +3,8 @@
 
 import pytest
 
-from genre.trie import Trie
 from genre.fairseq_model import GENRE, GENREHubInterface
+from genre.trie import Trie
 
 
 @pytest.fixture(scope="session")