Fix annotations

This commit is contained in:
Samuel Sloniker 2022-07-17 17:08:11 -07:00
parent 48639f5d8d
commit b36d8e6081

View File

@ -6,7 +6,7 @@ from typing import Iterable, Mapping, List, Dict, Union
def compile( def compile(
raw_model: Iterable[Mapping[str, str]], max_ngram_length: int = 1 raw_model: Iterable[Mapping[str, str]], max_ngram_length: int = 1
) -> Dict[str, Union[int, List[Union[str, int]]]]: ) -> Dict[str, Union[str, int, List[int], List[str]]]:
"""Compile a raw model. """Compile a raw model.
Parameters Parameters
@ -60,14 +60,15 @@ def compile(
except KeyError: except KeyError:
word_weights[word] = {category: value} word_weights[word] = {category: value}
model: Dict[str, Union[int, List[Union[str, int]]]] = {} model: Dict[str, Union[str, int, List[int], List[str]]] = {}
for word, weights in word_weights.items(): for word, weights in word_weights.items():
total = sum(weights.values()) total = sum(weights.values())
model[word] = [] new_weights: List[int] = []
for category in names: for category in names:
model[word].append( new_weights.append(
round((weights.get(category, 0) / total) * 65535) round((weights.get(category, 0) / total) * 65535)
) )
model[word] = new_weights
model["__names__"] = names model["__names__"] = names
model["__ngrams__"] = max_ngram_length model["__ngrams__"] = max_ngram_length