Skip to content

huggingface_provider

HuggingFaceProvider

Local HuggingFace embedding provider using sentence-transformers.

Source code in wintermute/ai/providers/huggingface_provider.py 
 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
class HuggingFaceProvider:
    """Local HuggingFace embedding provider using sentence-transformers."""

    def __init__(self, name: str = "local_embedder") -> None:
        self._name = name
        self._models: dict[str, Any] = {}
        if SentenceTransformer is None:
            log.warning(
                "sentence-transformers not installed. HuggingFaceProvider will not work."
            )

    @property
    def name(self) -> str:
        return self._name

    @property
    def description(self) -> str:
        return "Local HuggingFace sentence-transformers for embeddings only."

    def list_models(self) -> list[ModelInfo]:
        """Returns a list of common embedding models supported."""
        return [
            ModelInfo(
                name="all-MiniLM-L6-v2",
                family="BERT",
                context_window=512,
                supports_tools=False,
                supports_json=False,
                supports_stream=False,
            ),
            ModelInfo(
                name="BAAI/bge-small-en-v1.5",
                family="BERT",
                context_window=512,
                supports_tools=False,
                supports_json=False,
                supports_stream=False,
            ),
        ]

    def chat(self, req: ChatRequest) -> ChatResponse:
        raise NotImplementedError(
            "HuggingFace provider is currently for embeddings only."
        )

    def embed(
        self, texts: Iterable[str], model: Optional[str] = None
    ) -> list[list[float]]:
        """
        Embeds texts using a local SentenceTransformer model.
        Args:
            texts: List of strings to embed.
            model: Model name/path (default: 'all-MiniLM-L6-v2').
        """
        if SentenceTransformer is None:
            raise ImportError("sentence-transformers not installed.")

        model_name = model or "all-MiniLM-L6-v2"

        if model_name not in self._models:
            log.info(f"Loading local embedding model: {model_name}")
            self._models[model_name] = SentenceTransformer(model_name)

        transformer = self._models[model_name]
        # encode returns numpy array or list of numpy arrays; convert to list[list[float]]
        embeddings = transformer.encode(list(texts), convert_to_numpy=True)
        return cast("list[list[float]]", embeddings.tolist())

    def count_tokens(self, text: str, model: Optional[str] = None) -> int:
        """
        Returns estimated token count. For now, using a simple heuristic
        as exact tokenization requires the specific tokenizer.
        """
        if SentenceTransformer is None:
            return len(text.split())

        model_name = model or "all-MiniLM-L6-v2"
        # If model is loaded, we can use its tokenizer
        if model_name in self._models:
            tokenizer = self._models[model_name].tokenizer
            return len(tokenizer.tokenize(text))

        # Fallback heuristic
        return len(text.split())

count_tokens(text, model=None)

Returns estimated token count. For now, using a simple heuristic as exact tokenization requires the specific tokenizer.

Source code in wintermute/ai/providers/huggingface_provider.py 
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
def count_tokens(self, text: str, model: Optional[str] = None) -> int:
    """
    Returns estimated token count. For now, using a simple heuristic
    as exact tokenization requires the specific tokenizer.
    """
    if SentenceTransformer is None:
        return len(text.split())

    model_name = model or "all-MiniLM-L6-v2"
    # If model is loaded, we can use its tokenizer
    if model_name in self._models:
        tokenizer = self._models[model_name].tokenizer
        return len(tokenizer.tokenize(text))

    # Fallback heuristic
    return len(text.split())

embed(texts, model=None)

Embeds texts using a local SentenceTransformer model. Args: texts: List of strings to embed. model: Model name/path (default: 'all-MiniLM-L6-v2').

Source code in wintermute/ai/providers/huggingface_provider.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
def embed(
    self, texts: Iterable[str], model: Optional[str] = None
) -> list[list[float]]:
    """
    Embeds texts using a local SentenceTransformer model.
    Args:
        texts: List of strings to embed.
        model: Model name/path (default: 'all-MiniLM-L6-v2').
    """
    if SentenceTransformer is None:
        raise ImportError("sentence-transformers not installed.")

    model_name = model or "all-MiniLM-L6-v2"

    if model_name not in self._models:
        log.info(f"Loading local embedding model: {model_name}")
        self._models[model_name] = SentenceTransformer(model_name)

    transformer = self._models[model_name]
    # encode returns numpy array or list of numpy arrays; convert to list[list[float]]
    embeddings = transformer.encode(list(texts), convert_to_numpy=True)
    return cast("list[list[float]]", embeddings.tolist())

list_models()

Returns a list of common embedding models supported.

Source code in wintermute/ai/providers/huggingface_provider.py 
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def list_models(self) -> list[ModelInfo]:
    """Returns a list of common embedding models supported."""
    return [
        ModelInfo(
            name="all-MiniLM-L6-v2",
            family="BERT",
            context_window=512,
            supports_tools=False,
            supports_json=False,
            supports_stream=False,
        ),
        ModelInfo(
            name="BAAI/bge-small-en-v1.5",
            family="BERT",
            context_window=512,
            supports_tools=False,
            supports_json=False,
            supports_stream=False,
        ),
    ]

register(as_name='local_embedder')

Registers the HuggingFaceProvider.

Source code in wintermute/ai/providers/huggingface_provider.py 
128
129
130
131
132
133
134
135
def register(as_name: str = "local_embedder") -> None:
    """Registers the HuggingFaceProvider."""
    try:
        provider = HuggingFaceProvider(name=as_name)
        llms.register(provider)
        log.info(f"Registered HuggingFaceProvider as '{as_name}'")
    except Exception as e:
        log.error(f"Failed to register HuggingFaceProvider: {e}")