mirror of https://github.com/morpheus65535/bazarr
73 lines
1.9 KiB
Python
73 lines
1.9 KiB
Python
"""
|
|
🌏 Charamel: Truly Universal Encoding Detection in Python 🌎
|
|
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
|
|
|
Licensed under Apache 2.0
|
|
"""
|
|
import gzip
|
|
import pathlib
|
|
import struct
|
|
from typing import Any, Dict, List, Sequence
|
|
|
|
from charamel.encoding import Encoding
|
|
|
|
RESOURCE_DIRECTORY = pathlib.Path(__file__).parent.absolute()
|
|
WEIGHT_DIRECTORY = RESOURCE_DIRECTORY / 'weights'
|
|
|
|
|
|
def _unpack(file: pathlib.Path, pattern: str) -> List[Any]:
|
|
"""
|
|
Unpack struct values from file
|
|
|
|
Args:
|
|
file: File that stores struct-packed values
|
|
pattern: Struct pattern
|
|
|
|
Returns:
|
|
List of unpacked values
|
|
"""
|
|
with gzip.open(file, 'rb') as data:
|
|
return [values[0] for values in struct.iter_unpack(pattern, data.read())]
|
|
|
|
|
|
def load_features() -> Dict[int, int]:
|
|
"""
|
|
Load byte-level feature names and indices
|
|
|
|
Returns:
|
|
Mapping from features to their indices in weight matrix
|
|
"""
|
|
features = _unpack(RESOURCE_DIRECTORY / 'features.gzip', pattern='>H')
|
|
return {feature: index for index, feature in enumerate(features)}
|
|
|
|
|
|
def load_biases(encodings: Sequence[Encoding]) -> Dict[Encoding, float]:
|
|
"""
|
|
Load linear model bias values for given encodings
|
|
|
|
Args:
|
|
encodings: List of encodings
|
|
|
|
Returns:
|
|
Mapping from encodings to their biases
|
|
"""
|
|
biases = {}
|
|
with gzip.open(RESOURCE_DIRECTORY / 'biases.gzip', 'rb') as data:
|
|
for line in data:
|
|
encoding, bias = line.decode().split()
|
|
biases[encoding] = float(bias)
|
|
|
|
return {encoding: biases[encoding] for encoding in encodings}
|
|
|
|
|
|
def load_weights(encodings: Sequence[Encoding]) -> Dict[Encoding, List[float]]:
|
|
"""
|
|
|
|
:param encodings:
|
|
:return:
|
|
"""
|
|
weights = {}
|
|
for encoding in encodings:
|
|
weights[encoding] = _unpack(WEIGHT_DIRECTORY / f'{encoding}.gzip', pattern='>e')
|
|
return weights
|