Adding sentence attributes:

-relation_distribiution - probas of arc between dependent and head -relation_label_distribiution- probas for each label in sentence

combo/data/api.py

@@ -24,16 +24,21 @@ class Token:
     embeddings: Dict[str, List[float]] = field(default_factory=list, repr=False)
 
 
+
 @dataclass
 class Sentence:
     tokens: List[Token] = field(default_factory=list)
     sentence_embedding: List[float] = field(default_factory=list, repr=False)
+    relation_distribution: List[float] = field(default_factory=list, repr=False)
+    relation_label_distribution: List[float] = field(default_factory=list, repr=False)
     metadata: Dict[str, Any] = field(default_factory=collections.OrderedDict)
 
     def to_json(self):
         return json.dumps({
             "tokens": [dataclasses.asdict(t) for t in self.tokens],
             "sentence_embedding": self.sentence_embedding,
+            "relation_distribution": self.relation_distribution,
+            "relation_label_distribution": self.relation_label_distribution,
             "metadata": self.metadata,
         })
 
@@ -79,7 +84,12 @@ def tokens2conllu(tokens: List[str]) -> conllu.TokenList:
 
 
 def conllu2sentence(conllu_sentence: conllu.TokenList,
-                    sentence_embedding=None, embeddings=None) -> Sentence:
+                    sentence_embedding=None, embeddings=None, relation_distribution=None,
+                    relation_label_distribution=None) -> Sentence:
+    if relation_distribution is None:
+        relation_distribution = []
+    if relation_label_distribution is None:
+        relation_label_distribution = []
     if embeddings is None:
         embeddings = {}
     if sentence_embedding is None:
@@ -94,5 +104,7 @@ def conllu2sentence(conllu_sentence: conllu.TokenList,
     return Sentence(
         tokens=tokens,
         sentence_embedding=sentence_embedding,
+        relation_distribution=relation_distribution,
+        relation_label_distribution=relation_label_distribution,
         metadata=conllu_sentence.metadata
     )