tests/test_text_features.py

import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

import numpy as np

from bucklake_ai.config import Settings
from bucklake_ai.text_features import TextFeatureEncoder


class StubToken:
    def __init__(
        self, pos: str, tag: str, *, is_punct: bool = False, is_space: bool = False
    ):
        self.pos_ = pos
        self.tag_ = tag
        self.is_punct = is_punct
        self.is_space = is_space


class StubEntity:
    def __init__(self, label: str, text: str):
        self.label_ = label
        self.text = text


class StubDoc:
    def __init__(self):
        self.ents = [StubEntity("ORG", "Apple")]
        self._tokens = [
            StubToken("PROPN", "NNP"),
            StubToken("VERB", "VBZ"),
            StubToken("PUNCT", ".", is_punct=True),
        ]

    def __iter__(self):
        return iter(self._tokens)


class TextFeatureEncoderTests(unittest.TestCase):
    def _settings(self) -> Settings:
        root_dir = Path(tempfile.mkdtemp())
        return Settings(
            app_name="BuckLakeAI",
            app_version="2.0.0",
            model_version="anchored-path-v1",
            input_contract_version="v2.1.0",
            root_dir=root_dir,
            hf_repo_id="parkerjj/BuckLake-Stock-Model",
            hf_model_filename="stock_prediction_model_anchored-path-v2_best_round45.keras",
            hf_scaler_filename="scalers_v21_rolling_7d.json",
            hf_cache_dir=root_dir / ".cache" / "huggingface",
            hf_token=None,
            model_weights_path=root_dir
            / ".cache"
            / "huggingface"
            / "stock_prediction_model_anchored-path-v2_best_round45.keras",
            text_encoder_model="test-encoder",
            text_encoder_device="cpu",
            text_encoder_max_seq_length=512,
            enable_finbert_sentiment=False,
            finbert_model_name="ProsusAI/finbert",
            preload_model=False,
            preload_text_encoder=False,
            scaler_artifact_path=root_dir / "data" / "scalers_v21_rolling_7d.json",
        )

    def test_build_pos_entity_vectors_matches_training_hash_semantics(self):
        encoder = TextFeatureEncoder(self._settings())

        with patch.object(encoder, "_load_spacy_nlp") as load_spacy:
            load_spacy.return_value.return_value = StubDoc()
            pos_vector, entity_vector = encoder.build_pos_entity_vectors(
                "Apple launches something new."
            )

        self.assertEqual(pos_vector.shape, (1024,))
        self.assertEqual(entity_vector.shape, (1024,))
        self.assertEqual(pos_vector.dtype, np.float32)
        self.assertEqual(entity_vector.dtype, np.float32)
        self.assertGreater(float(np.linalg.norm(pos_vector)), 0.0)
        self.assertGreater(float(np.linalg.norm(entity_vector)), 0.0)
        self.assertAlmostEqual(float(pos_vector.sum()), 1.0)
        self.assertAlmostEqual(float(entity_vector.sum()), 1.0)


if __name__ == "__main__":
    unittest.main()