Downloads: 1
Diff of
/annotations.py
[000000]
..
[1de6ed]
Switch to unified view
| a | b/annotations.py | ||
|---|---|---|---|
| 1 | from typing import List, Tuple, Optional |
||
| 2 | |||
| 3 | |||
| 4 | class Annotation: |
||
| 5 | """ |
||
| 6 | A generic class for annotations |
||
| 7 | """ |
||
| 8 | |||
| 9 | def __init__(self, ann_id: str, ann_name: str = None) -> None: |
||
| 10 | self.ann_id = ann_id |
||
| 11 | self.name = ann_name |
||
| 12 | |||
| 13 | |||
| 14 | class Entity(Annotation): |
||
| 15 | """ |
||
| 16 | Objects that represent named entities |
||
| 17 | """ |
||
| 18 | |||
| 19 | def __init__(self, entity_id: str, |
||
| 20 | entity_type: str = None, |
||
| 21 | char_range: List[int] = None) -> None: |
||
| 22 | """ |
||
| 23 | Initializes Entity object. |
||
| 24 | |||
| 25 | Parameters |
||
| 26 | ---------- |
||
| 27 | entity_id : str |
||
| 28 | Unique entity ID. |
||
| 29 | entity_type : str, optional |
||
| 30 | The type of entity. The default is None. |
||
| 31 | |||
| 32 | """ |
||
| 33 | super().__init__(entity_id, entity_type) |
||
| 34 | if char_range is None: |
||
| 35 | self.range = [None, None] |
||
| 36 | else: |
||
| 37 | self.range: List[int] = char_range |
||
| 38 | self.ann_text: str = "" |
||
| 39 | self.relation_group: str = "" |
||
| 40 | |||
| 41 | def set_range(self, new_range: List[int]) -> None: |
||
| 42 | """ |
||
| 43 | Add annotation range |
||
| 44 | """ |
||
| 45 | self.range = new_range |
||
| 46 | |||
| 47 | def set_text(self, text: str) -> None: |
||
| 48 | """ |
||
| 49 | Sets the annotation text |
||
| 50 | """ |
||
| 51 | self.ann_text = text |
||
| 52 | |||
| 53 | def set_entity_type(self, entity_type: str) -> None: |
||
| 54 | """ |
||
| 55 | Sets the entity type |
||
| 56 | """ |
||
| 57 | self.name = entity_type |
||
| 58 | |||
| 59 | def __repr__(self) -> str: |
||
| 60 | """ |
||
| 61 | String representation of the object |
||
| 62 | """ |
||
| 63 | string = "\n" |
||
| 64 | string += "ID: " + self.ann_id + "\n" |
||
| 65 | string += "Entity name: " + str(self.name) + "\n" |
||
| 66 | |||
| 67 | string += "Character range: " |
||
| 68 | string += str(self.range[0]) + " " + str(self.range[1]) + "\n" |
||
| 69 | |||
| 70 | if self.ann_text: |
||
| 71 | string += "Entity text: '" + str(self.ann_text) + "'" |
||
| 72 | |||
| 73 | return string |
||
| 74 | |||
| 75 | def __str__(self) -> str: |
||
| 76 | return self.__repr__() |
||
| 77 | |||
| 78 | def __getitem__(self, key: int) -> int: |
||
| 79 | """ |
||
| 80 | Makes the class subsriptable on range |
||
| 81 | """ |
||
| 82 | return self.range[key] |
||
| 83 | |||
| 84 | def __iter__(self) -> Tuple[int, int]: |
||
| 85 | """ |
||
| 86 | Makes class iterable on range |
||
| 87 | """ |
||
| 88 | yield self.range[0] |
||
| 89 | yield self.range[1] |
||
| 90 | |||
| 91 | def __eq__(self, other) -> bool: |
||
| 92 | """ |
||
| 93 | Overrides equality method |
||
| 94 | """ |
||
| 95 | if self.name == other.name and self.range == other.range: |
||
| 96 | return True |
||
| 97 | else: |
||
| 98 | return False |
||
| 99 | |||
| 100 | |||
| 101 | class Relation(Annotation): |
||
| 102 | """ |
||
| 103 | Objects that represent relations between named entities |
||
| 104 | """ |
||
| 105 | |||
| 106 | def __init__(self, relation_id: str, relation_type: str, |
||
| 107 | arg1: Entity = None, arg2: Entity = None) -> None: |
||
| 108 | |||
| 109 | super().__init__(relation_id, relation_type) |
||
| 110 | self.arg1 = arg1 |
||
| 111 | self.arg2 = arg2 |
||
| 112 | |||
| 113 | def set_entity_relation(self, arg1: str, arg2: str) -> None: |
||
| 114 | """ |
||
| 115 | Sets the entities that are related |
||
| 116 | """ |
||
| 117 | self.arg1 = arg1 |
||
| 118 | self.arg2 = arg2 |
||
| 119 | |||
| 120 | def get_entities(self) -> List[Optional[Entity]]: |
||
| 121 | """ |
||
| 122 | Returns related entities |
||
| 123 | """ |
||
| 124 | return [self.arg1, self.arg2] |
||
| 125 | |||
| 126 | def set_relation_type(self, relation_type: str) -> None: |
||
| 127 | """ |
||
| 128 | Sets the relation type |
||
| 129 | """ |
||
| 130 | self.name = relation_type |
||
| 131 | |||
| 132 | def __repr__(self) -> str: |
||
| 133 | """ |
||
| 134 | String representation of the object |
||
| 135 | """ |
||
| 136 | string = "\n" |
||
| 137 | string += "ID: " + str(self.ann_id) + "\n" |
||
| 138 | string += "Relation type: " + str(self.name) + "\n" |
||
| 139 | string += "\nEntity 1: \n" |
||
| 140 | string += self.arg1.__repr__() + "\n" |
||
| 141 | string += "\nEntity 2: \n" |
||
| 142 | string += self.arg2.__repr__() |
||
| 143 | |||
| 144 | return string |
||
| 145 | |||
| 146 | def __str__(self) -> str: |
||
| 147 | return self.__repr__() |
||
| 148 | |||
| 149 | def __eq__(self, other) -> bool: |
||
| 150 | """ |
||
| 151 | Overrides the default equality method. |
||
| 152 | """ |
||
| 153 | if self.arg1 == other.arg1 and self.arg2 == other.arg2: |
||
| 154 | return True |
||
| 155 | |||
| 156 | elif self.arg2 == other.arg1 and self.arg1 == other.arg2: |
||
| 157 | return True |
||
| 158 | |||
| 159 | else: |
||
| 160 | return False |