diff --git a/src/pastel/pastel.py b/src/pastel/pastel.py index 55dd1dc..625aead 100644 --- a/src/pastel/pastel.py +++ b/src/pastel/pastel.py @@ -95,15 +95,11 @@ def from_feature_list(feature_names: Sequence[FEATURE_TYPE]) -> "Pastel": return Pastel(new_model) @staticmethod - def load_model(model_file: str) -> "Pastel": - """Load model from JSON file. Convert any functions in the model - from their names to Callable functions.""" - - with open(model_file, "rt", encoding="utf-8") as json_in: - model_json = json.load(json_in) + def from_dict(model_dict: dict[str, float]) -> "Pastel": + "Create model from a map of features to weights" # replace function names with function objects found in pastel_functions module new_model = {} - for feature, weight in model_json.items(): + for feature, weight in model_dict.items(): if feature in pastel_functions.__all__: new_model[getattr(pastel_functions, feature)] = weight elif feature == "bias": @@ -113,6 +109,15 @@ def load_model(model_file: str) -> "Pastel": return Pastel(new_model) + @staticmethod + def load_model(model_file: str) -> "Pastel": + """Load model from JSON file. Convert any functions in the model + from their names to Callable functions.""" + + with open(model_file, "rt", encoding="utf-8") as json_in: + model_json = json.load(json_in) + return Pastel.from_dict(model_json) + def save_model(self, model_path: str) -> None: """ Save the questions, functions and associated weights to a local JSON file @@ -250,7 +255,7 @@ async def get_answers_to_questions( ) -> dict[Sentence, dict[FEATURE_TYPE, float]]: """Embed each example into the prompt and pass to genAI, then get answers for non-genAI functions. - For each sentence, this Returns a dictionary mapping features to scores.""" + For each sentence, this returns a dictionary mapping features to scores.""" jobs = [ self._get_answers_for_single_sentence(sentence) for sentence in sentences diff --git a/tests/pastel/test_pastel.py b/tests/pastel/test_pastel.py index ad82189..e6b96ed 100644 --- a/tests/pastel/test_pastel.py +++ b/tests/pastel/test_pastel.py @@ -36,6 +36,14 @@ def test_load_file(pastel_instance: Pastel) -> None: assert loaded.model == pastel_instance.model +def test_model_from_dict(): + model_dict = {"bias": 1.0, "question_0?": 2.0, "is_claim_type_personal": 3.0} + model = Pastel.from_dict(model_dict) + assert model.get_bias() == 1.0 + assert len(model.get_functions()) == 1 + assert len(model.get_questions()) == 1 + + def test_make_prompt(pastel_instance: Pastel) -> None: sentence = Sentence("The sky is blue.", tuple("quantity")) prompt = pastel_instance.make_prompt(sentence)