mirror of
https://github.com/open-webui/open-webui.git
synced 2026-10-22 11:23:05 -05:00
Classic298
a915df8ce1
The Docker image now runs Python 3.12 on Debian 13 (trixie) instead of Python 3.11 on Debian 12, and the pip install docs recommend 3.12. Image size stays about the same: the slim image is 1 MB smaller and the standard image is about 50 MB larger, mostly from the newer ffmpeg and pandoc that Debian 13 ships. Admins should know that custom functions, tools and anything installed with pip inside the container now run on Python 3.12, so a package that only has builds for 3.11 needs a newer version. Installing with pip still works on both Python 3.11 and 3.12.
92 lines
3.8 KiB
Python
92 lines
3.8 KiB
Python
import logging
|
|
import os
|
|
|
|
import numpy as np
|
|
import torch
|
|
from colbert.infra import ColBERTConfig
|
|
from colbert.modeling import base_colbert, hf_colbert
|
|
from colbert.modeling.checkpoint import Checkpoint
|
|
from open_webui.retrieval.models.base_reranker import BaseReranker
|
|
from transformers import PretrainedConfig
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
def class_factory(name_or_path: str) -> type:
|
|
hf_colbert_class = hf_colbert.class_factory(name_or_path)
|
|
|
|
# colbert-ai never calls post_init, which transformers 5 needs to finish setting up the model
|
|
class HF_ColBERT(hf_colbert_class):
|
|
def __init__(self, config: PretrainedConfig, colbert_config: ColBERTConfig) -> None:
|
|
super().__init__(config, colbert_config)
|
|
self.post_init()
|
|
|
|
return HF_ColBERT
|
|
|
|
|
|
base_colbert.class_factory = class_factory
|
|
|
|
|
|
class ColBERT(BaseReranker):
|
|
def __init__(self, name, **kwargs) -> None:
|
|
log.info('ColBERT: Loading model %s', name)
|
|
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
|
|
DOCKER = kwargs.get('env') == 'docker'
|
|
if DOCKER:
|
|
# This is a workaround for the issue with the docker container
|
|
# where the torch extension is not loaded properly
|
|
# and the following error is thrown:
|
|
# /root/.cache/torch_extensions/py312_cpu/segmented_maxsim_cpp/segmented_maxsim_cpp.so: cannot open shared object file: No such file or directory
|
|
|
|
lock_file = '/root/.cache/torch_extensions/py312_cpu/segmented_maxsim_cpp/lock'
|
|
if os.path.exists(lock_file):
|
|
os.remove(lock_file)
|
|
|
|
self.ckpt = Checkpoint(
|
|
name,
|
|
colbert_config=ColBERTConfig(model_name=name),
|
|
).to(self.device)
|
|
pass
|
|
|
|
def calculate_similarity_scores(self, query_embeddings, document_embeddings):
|
|
query_embeddings = query_embeddings.to(self.device)
|
|
document_embeddings = document_embeddings.to(self.device)
|
|
|
|
# Validate dimensions to ensure compatibility
|
|
if query_embeddings.dim() != 3:
|
|
raise ValueError(f'Expected query embeddings to have 3 dimensions, but got {query_embeddings.dim()}.')
|
|
if document_embeddings.dim() != 3:
|
|
raise ValueError(f'Expected document embeddings to have 3 dimensions, but got {document_embeddings.dim()}.')
|
|
if query_embeddings.size(0) not in [1, document_embeddings.size(0)]:
|
|
raise ValueError('There should be either one query or queries equal to the number of documents.')
|
|
|
|
# Transpose the query embeddings to align for matrix multiplication
|
|
transposed_query_embeddings = query_embeddings.permute(0, 2, 1)
|
|
# Compute similarity scores using batch matrix multiplication
|
|
computed_scores = torch.matmul(document_embeddings, transposed_query_embeddings)
|
|
# Apply max pooling to extract the highest semantic similarity across each document's sequence
|
|
maximum_scores = torch.max(computed_scores, dim=1).values
|
|
|
|
# Sum up the maximum scores across features to get the overall document relevance scores
|
|
final_scores = maximum_scores.sum(dim=1)
|
|
|
|
normalized_scores = torch.softmax(final_scores, dim=0)
|
|
|
|
return normalized_scores.detach().cpu().numpy().astype(np.float32)
|
|
|
|
def predict(self, sentences, batch_size=32):
|
|
query = sentences[0][0]
|
|
docs = [i[1] for i in sentences]
|
|
|
|
# Embedding the documents
|
|
embedded_docs = self.ckpt.docFromText(docs, bsize=batch_size)[0]
|
|
# Embedding the queries
|
|
embedded_queries = self.ckpt.queryFromText([query], bsize=batch_size)
|
|
embedded_query = embedded_queries[0]
|
|
|
|
# Calculate retrieval scores for the query against all documents
|
|
scores = self.calculate_similarity_scores(embedded_query.unsqueeze(0), embedded_docs)
|
|
|
|
return scores
|