Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

fix: Loading of embedddings #22260

Merged
merged 3 commits into from
May 20, 2024
Merged

fix: Loading of embedddings #22260

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
2809a48
Fix embeddings loading
benjackwhite May 11, 2024
005362c
Fix
benjackwhite May 11, 2024
bf38c12
Merge branch 'master' into fix/local-dev-embeddings
daibhin May 20, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 14 additions & 5 deletions ee/session_recordings/ai/embeddings_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import datetime
import pytz

from typing import Any
from typing import Any, Optional

from abc import ABC, abstractmethod
from prometheus_client import Histogram, Counter
Expand All @@ -23,9 +23,18 @@
only_pageview_urls,
)

# tiktoken.encoding_for_model(model_name) specifies encoder
# model_name = "text-embedding-3-small" for this usecase
encoding = tiktoken.get_encoding("cl100k_base")
_encoding: Optional[tiktoken.Encoding] = None


def get_encoding() -> tiktoken.Encoding:
global _encoding
if not _encoding:
# NOTE: This does an API request so we want to ensure we load it lazily and not at startup
# tiktoken.encoding_for_model(model_name) specifies encoder
# model_name = "text-embedding-3-small" for this usecase
_encoding = tiktoken.get_encoding("cl100k_base")
return _encoding
Comment on lines +29 to +36
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The implementation of get_encoding supports lazy loading effectively. Consider adding error handling for the API request to tiktoken.get_encoding.

def get_encoding() -> tiktoken.Encoding:
    global _encoding
    if not _encoding:
        try:
            _encoding = tiktoken.get_encoding("cl100k_base")
        except Exception as e:
            logger.error("Failed to load encoding", error=str(e))
            raise
    return _encoding



MAX_TOKENS_FOR_MODEL = 8191

Expand Down Expand Up @@ -194,7 +203,7 @@ def _embed(self, input: str, source_type: str):

def _num_tokens_for_input(self, string: str) -> int:
"""Returns the number of tokens in a text string."""
return len(encoding.encode(string))
return len(get_encoding().encode(string))

def _flush_embeddings_to_clickhouse(self, embeddings: list[dict[str, Any]], source_type: str) -> None:
try:
Expand Down
Loading