File size: 3,170 Bytes
96b6673 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
import numpy as np
from numpy.typing import NDArray
from typing import Optional, List
from abc import ABC, abstractmethod
from .utils import split_text
class BaseContextPartitioner(ABC):
"""
A base class for partitioning a context into sources.
Attributes:
context (str):
The context to partition.
Methods:
num_sources(self) -> int:
Property. The number of sources within the context.
split_context(self) -> None:
Split the context into sources.
get_source(self, index: int) -> str:
Get a represention of the source corresponding to a given index.
get_context(self, mask: Optional[NDArray] = None) -> str:
Get a version of the context ablated according to the given mask.
sources(self) -> List[str]:
Property. A list of all sources within the context.
"""
def __init__(self, context: str) -> None:
self.context = context
@property
@abstractmethod
def num_sources(self) -> int:
"""The number of sources."""
@abstractmethod
def split_context(self) -> None:
"""Split the context into sources."""
@abstractmethod
def get_source(self, index: int) -> str:
"""Get a represention of the source corresponding to a given index."""
@abstractmethod
def get_context(self, mask: Optional[NDArray] = None):
"""Get a version of the context ablated according to the given mask."""
@property
def sources(self) -> List[str]:
"""A list of all sources."""
return [self.get_source(i) for i in range(self.num_sources)]
class SimpleContextPartitioner(BaseContextPartitioner):
"""
A simple context partitioner that splits the context into sources based on
a separator.
"""
def __init__(self, context: str, source_type: str = "sentence") -> None:
super().__init__(context)
self.source_type = source_type
self._cache = {}
def split_context(self):
"""Split text into parts and cache the parts and separators."""
parts, separators, _ = split_text(self.context, self.source_type)
self._cache["parts"] = parts
self._cache["separators"] = separators
@property
def parts(self):
if self._cache.get("parts") is None:
self.split_context()
return self._cache["parts"]
@property
def separators(self):
if self._cache.get("separators") is None:
self.split_context()
return self._cache["separators"]
@property
def num_sources(self) -> int:
return len(self.parts)
def get_source(self, index: int) -> str:
return self.parts[index]
def get_context(self, mask: Optional[NDArray] = None):
if mask is None:
mask = np.ones(self.num_sources, dtype=bool)
separators = np.array(self.separators)[mask]
parts = np.array(self.parts)[mask]
context = ""
for i, (separator, part) in enumerate(zip(separators, parts)):
if i > 0:
context += separator
context += part
return context
|