added role list to labels

Riccorl · Dec 2, 2020 · 0bded03 · 0bded03
1 parent 95ee1cd
commit 0bded03
Show file tree

Hide file tree

Showing 3 changed files with 39,903 additions and 4 deletions.
diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="transformer_srl",  # Replace with your own username
-    version="2.4.9",
+    version="2.4.10",
     author="Riccardo Orlando",
     author_email="orlandoricc@gmail.com",
     description="SRL Transformer model",

diff --git a/transformer_srl/models.py b/transformer_srl/models.py
@@ -22,7 +22,7 @@
 from transformer_srl.utils import load_label_list, load_lemma_frame, load_role_frame
 
 LEMMA_FRAME_PATH = pathlib.Path(__file__).resolve().parent / "resources" / "lemma2va_ml.tsv"
-FRAME_ROLE_PATH = pathlib.Path(__file__).resolve().parent / "resources" / "frame2role.csv"
+FRAME_ROLE_PATH = pathlib.Path(__file__).resolve().parent / "resources" / "frame2role_ml.csv"
 FRAME_LIST_PATH = pathlib.Path(__file__).resolve().parent / "resources" / "framelist.txt"
 ROLE_LIST_PATH = pathlib.Path(__file__).resolve().parent / "resources" / "rolelist.txt"
 
@@ -66,8 +66,8 @@ def __init__(
         Model.__init__(self, vocab, **kwargs)
         self.lemma_frame_dict = load_lemma_frame(LEMMA_FRAME_PATH)
         self.frame_role_dict = load_role_frame(FRAME_ROLE_PATH)
-        self.restrict_frames = True
-        self.restrict_roles = True
+        self.restrict_frames = restrict_frames
+        self.restrict_roles = restrict_roles
         self.transformer = AutoModel.from_pretrained(bert_model)
         self.frame_criterion = nn.CrossEntropyLoss()
         if inventory == "verbatlas":