Spaces:
Sleeping
Sleeping
File size: 11,248 Bytes
a54024a 27a1559 a54024a 27a1559 a54024a 27a1559 a54024a |
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 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 |
import string
from typing import List, Union
import nltk
import torch
from numpy.typing import NDArray
from .type_aliases import DEVICE_TYPE, ENCODER_DEVICE_TYPE, NumSentencesType, EmbeddingSlicesType
def get_gpu(gpu: DEVICE_TYPE) -> ENCODER_DEVICE_TYPE:
"""
Determine the correct GPU device based on the provided input. In the following, output 0 means CUDA device 0.
Args:
gpu (Union[bool, str, int, List[Union[str, int]]]): Input specifying the GPU device(s):
- bool: If True, returns 0 if CUDA is available, otherwise returns "cpu".
- str: Can be "cpu", "gpu", or "cuda" (case-insensitive). Returns 0 if CUDA is available
and the input is not "cpu", otherwise returns "cpu".
- int: Should be a valid GPU index. Returns the index if CUDA is available and valid,
otherwise returns "cpu".
- List[Union[str, int]]: List containing combinations of the str/int. Processes each
element and returns a list of corresponding results.
Returns:
Union[str, int, List[Union[str, int]]]: Depending on the input type:
- str: Returns "cpu" if no GPU is available or the input is "cpu".
- int: Returns the GPU index if valid and CUDA is available.
- List[Union[str, int]]: Returns a list of strings and/or integers based on the input list.
Raises:
ValueError: If the input gpu type is not recognized or invalid.
ValueError: If a string input is not one of ["cpu", "gpu", "cuda"].
ValueError: If an integer input is outside the valid range of GPU indices.
Notes:
- This function checks CUDA availability using torch.cuda.is_available() and counts
available GPUs using torch.cuda.device_count().
- Case insensitivity is maintained for string inputs ("cpu", "gpu", "cuda").
- The function ensures robust error handling for invalid input types or out-of-range indices.
"""
# Ensure gpu index is within the range of total available gpus
gpu_available = torch.cuda.is_available()
gpu_count = torch.cuda.device_count()
correct_strs = ["cpu", "gpu", "cuda"]
def _get_single_device(gpu_item):
if isinstance(gpu_item, bool):
return 0 if gpu_item and gpu_available else "cpu"
elif isinstance(gpu_item, str):
if gpu_item.lower() not in correct_strs:
raise ValueError(f"Wrong gpu type: {gpu_item}. Should be one of {correct_strs}")
return 0 if (gpu_item.lower() != "cpu") and gpu_available else "cpu"
elif isinstance(gpu_item, int):
if gpu_item >= gpu_count:
raise ValueError(
f"There are {gpu_count} GPUs available. Provide a valid GPU index. You provided: {gpu_item}"
)
return gpu_item if gpu_available else "cpu"
else:
raise ValueError(f"Invalid gpu type: {type(gpu_item)}. Must be bool, str, or int.")
if isinstance(gpu, list):
seen_indices = set()
result = []
for item in gpu:
device = _get_single_device(item)
if isinstance(device, int):
if device not in seen_indices:
seen_indices.add(device)
result.append(device)
else:
result.append(device)
return result[0] if len(result) == 1 else result
else:
return _get_single_device(gpu)
def prep_sentences(sentences: List[str]) -> List[str]:
"""
Processes a list of sentences by stripping whitespace (at beginning and the end),
, filtering out empty sentences or sentences that only contains punctuations.
Args:
sentences (List[str]): A list of sentences to be processed.
Returns:
List[str]: A list of cleaned sentences
Raises:
ValueError: If the resulting list of sentences is empty.
Example:
>>> prep_sentences(["Hello, world!", " This is a test. ", "!!!"])
['Hello, world!', 'This is a test.']
>>> prep_sentences(["!!!", "..."])
ValueError: Document can't be empty.
"""
out = []
for sent in sentences:
sent = sent.strip()
sent_wo_punctuation = (
sent.translate(str.maketrans("", "", string.punctuation))
).strip()
if sent_wo_punctuation:
out.append(sent)
if len(out) == 0:
raise ValueError("Document can't be empty.")
return out
def tokenize_and_prep_document(document: Union[str, List[str]], tokenize: bool) -> List[str]:
"""
Tokenizes and prepares a document by either tokenizing it into sentences and processing each sentence,
or directly processing each element if `tokenize` is False.
Args:
document (Union[str, List[str]]): The document to be processed. It can be a single string (enitre document) or a
list of strings (list of sentences).
tokenize (bool): If True, tokenizes `document` into sentences using NLTK's sentence tokenizer before processing.
If False, processes each element of `document` directly as sentences.
Returns:
List[str]: A list of cleaned sentences.
Raises:
ValueError: If the resulting list of sentences is empty after processing.
Example:
>>> tokenize_and_prep_document("Hello, world! This is a test.", True)
['Hello, world!', 'This is a test.']
>>> tokenize_and_prep_document(["Hello, world!", "This is a test."], False)
['Hello, world!', 'This is a test.']
>>> tokenize_and_prep_document("!!! ...", True)
ValueError: Document can't be empty.
Note: Only the following two cases are possible.
tokenizer=True -> document: str
tokenizer=False -> document: List[str].
"""
if tokenize:
return prep_sentences(nltk.tokenize.sent_tokenize(document))
return prep_sentences(document)
def flatten_list(nested_list: list) -> list:
"""
Recursively flattens a nested list of any depth.
Parameters:
nested_list (list): The nested list to flatten.
Returns:
list: A flat list containing all the elements of the nested list.
"""
flat_list = []
for item in nested_list:
if isinstance(item, list):
flat_list.extend(flatten_list(item))
else:
flat_list.append(item)
return flat_list
def is_nested_list_of_type(lst_obj, element_type, depth: int) -> bool:
"""
Check if the given object is a nested list of a specific type up to a specified depth.
Args:
- lst_obj: The object to check, expected to be a list or a single element.
- element_type: The type that each element in the nested list should match.
- depth (int): The depth of nesting to check. Must be non-negative.
Returns:
- bool: True if lst_obj is a nested list of the specified type up to the given depth, False otherwise.
Raises:
- ValueError: If depth is negative.
Example:
```python
# Test cases
is_nested_list_of_type("test", str, 0) # Returns True
is_nested_list_of_type([1, 2, 3], str, 0) # Returns False
is_nested_list_of_type(["apple", "banana"], str, 1) # Returns True
is_nested_list_of_type([[1, 2], [3, 4]], int, 2) # Returns True
is_nested_list_of_type([[1, 2], ["a", "b"]], int, 2) # Returns False
is_nested_list_of_type([[[1], [2]], [[3], [4]]], int, 3) # Returns True
```
Explanation:
- The function checks if `lst_obj` is a nested list of elements of type `element_type` up to `depth` levels deep.
- If `depth` is 0, it checks if `lst_obj` itself is of type `element_type`.
- If `depth` is greater than 0, it recursively checks each level of nesting to ensure all elements match `element_type`.
- Raises a `ValueError` if `depth` is negative, as depth must be a non-negative integer.
"""
if depth == 0:
return isinstance(lst_obj, element_type)
elif depth > 0:
return isinstance(lst_obj, list) and all(
is_nested_list_of_type(item, element_type, depth - 1) for item in lst_obj)
else:
raise ValueError("Depth can't be negative")
def slice_embeddings(embeddings: NDArray, num_sentences: NumSentencesType) -> EmbeddingSlicesType:
"""
Slice embeddings into segments based on the provided number of sentences per segment.
Args:
- embeddings (np.ndarray): The array of embeddings to be sliced.
- num_sentences (Union[List[int], List[List[int]]]):
- If a list of integers: Specifies the number of embeddings to take in each slice.
- If a list of lists of integers: Specifies multiple nested levels of slicing.
Returns:
- List[np.ndarray]: A list of numpy arrays where each array represents a slice of embeddings.
Raises:
- TypeError: If `num_sentences` is not of type List[int] or List[List[int]].
Example Usage:
```python
embeddings = np.random.rand(10, 5)
num_sentences = [3, 2, 5]
result = slice_embeddings(embeddings, num_sentences)
# `result` will be a list of numpy arrays:
# [embeddings[:3], embeddings[3:5], embeddings[5:]]
num_sentences_nested = [[2, 1], [3, 4]]
result_nested = slice_embeddings(embeddings, num_sentences_nested)
# `result_nested` will be a nested list of numpy arrays:
# [[embeddings[:2], embeddings[2:3]], [embeddings[3:6], embeddings[6:]]]
slice_embeddings(embeddings, "invalid") # Raises a TypeError
```
"""
def _slice_embeddings(s_idx: int, n_sentences: List[int]):
"""
Helper function to slice embeddings starting from index `s_idx`.
Args:
- s_idx (int): Starting index for slicing.
- n_sentences (List[int]): List specifying number of sentences in each slice.
Returns:
- Tuple[List[np.ndarray], int]: A tuple containing a list of sliced embeddings and the next starting index.
"""
_result = []
for count in n_sentences:
_result.append(embeddings[s_idx:s_idx + count])
s_idx += count
return _result, s_idx
if isinstance(num_sentences, list) and all(isinstance(item, int) for item in num_sentences):
result, _ = _slice_embeddings(0, num_sentences)
return result
elif isinstance(num_sentences, list) and all(
isinstance(sublist, list) and all(
isinstance(item, int) for item in sublist
)
for sublist in num_sentences
):
nested_result = []
start_idx = 0
for nested_num_sentences in num_sentences:
embedding_slice, start_idx = _slice_embeddings(start_idx, nested_num_sentences)
nested_result.append(embedding_slice)
return nested_result
else:
raise TypeError(f"Incorrect Type for {num_sentences=}")
|