Spaces:
Runtime error
Runtime error
File size: 720 Bytes
35b22df |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 |
"""Mock embedding model."""
from typing import Any, List
from gpt_index.embeddings.base import BaseEmbedding
class MockEmbedding(BaseEmbedding):
"""Mock embedding.
Used for token prediction.
Args:
embed_dim (int): embedding dimension
"""
def __init__(self, embed_dim: int, *args: Any, **kwargs: Any) -> None:
"""Init params."""
super().__init__(*args, **kwargs)
self.embed_dim = embed_dim
def _get_query_embedding(self, query: str) -> List[float]:
"""Get query embedding."""
return [0.5] * self.embed_dim
def _get_text_embedding(self, text: str) -> List[float]:
"""Get text embedding."""
return [0.5] * self.embed_dim
|