refac
This commit is contained in:
Timothy Jaeryang Baek
@@ -13,19 +13,17 @@ log = logging.getLogger(__name__)
|
||||
|
||||
class ColBERT(BaseReranker):
|
||||
def __init__(self, name, **kwargs) -> None:
|
||||
log.info("ColBERT: Loading model", name)
|
||||
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
log.info('ColBERT: Loading model', name)
|
||||
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||
|
||||
DOCKER = kwargs.get("env") == "docker"
|
||||
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/py311_cpu/segmented_maxsim_cpp/segmented_maxsim_cpp.so: cannot open shared object file: No such file or directory
|
||||
|
||||
lock_file = (
|
||||
"/root/.cache/torch_extensions/py311_cpu/segmented_maxsim_cpp/lock"
|
||||
)
|
||||
lock_file = '/root/.cache/torch_extensions/py311_cpu/segmented_maxsim_cpp/lock'
|
||||
if os.path.exists(lock_file):
|
||||
os.remove(lock_file)
|
||||
|
||||
@@ -36,23 +34,16 @@ class ColBERT(BaseReranker):
|
||||
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()}."
|
||||
)
|
||||
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()}."
|
||||
)
|
||||
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."
|
||||
)
|
||||
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)
|
||||
@@ -69,7 +60,6 @@ class ColBERT(BaseReranker):
|
||||
return normalized_scores.detach().cpu().numpy().astype(np.float32)
|
||||
|
||||
def predict(self, sentences):
|
||||
|
||||
query = sentences[0][0]
|
||||
docs = [i[1] for i in sentences]
|
||||
|
||||
@@ -80,8 +70,6 @@ class ColBERT(BaseReranker):
|
||||
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
|
||||
)
|
||||
scores = self.calculate_similarity_scores(embedded_query.unsqueeze(0), embedded_docs)
|
||||
|
||||
return scores
|
||||
|
||||
@@ -15,8 +15,8 @@ class ExternalReranker(BaseReranker):
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
url: str = "http://localhost:8080/v1/rerank",
|
||||
model: str = "reranker",
|
||||
url: str = 'http://localhost:8080/v1/rerank',
|
||||
model: str = 'reranker',
|
||||
timeout: Optional[int] = None,
|
||||
):
|
||||
self.api_key = api_key
|
||||
@@ -24,33 +24,31 @@ class ExternalReranker(BaseReranker):
|
||||
self.model = model
|
||||
self.timeout = timeout
|
||||
|
||||
def predict(
|
||||
self, sentences: List[Tuple[str, str]], user=None
|
||||
) -> Optional[List[float]]:
|
||||
def predict(self, sentences: List[Tuple[str, str]], user=None) -> Optional[List[float]]:
|
||||
query = sentences[0][0]
|
||||
docs = [i[1] for i in sentences]
|
||||
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"query": query,
|
||||
"documents": docs,
|
||||
"top_n": len(docs),
|
||||
'model': self.model,
|
||||
'query': query,
|
||||
'documents': docs,
|
||||
'top_n': len(docs),
|
||||
}
|
||||
|
||||
try:
|
||||
log.info(f"ExternalReranker:predict:model {self.model}")
|
||||
log.info(f"ExternalReranker:predict:query {query}")
|
||||
log.info(f'ExternalReranker:predict:model {self.model}')
|
||||
log.info(f'ExternalReranker:predict:query {query}')
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
'Content-Type': 'application/json',
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
}
|
||||
|
||||
if ENABLE_FORWARD_USER_INFO_HEADERS and user:
|
||||
headers = include_user_info_headers(headers, user)
|
||||
|
||||
r = requests.post(
|
||||
f"{self.url}",
|
||||
f'{self.url}',
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=self.timeout,
|
||||
@@ -60,13 +58,13 @@ class ExternalReranker(BaseReranker):
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
|
||||
if "results" in data:
|
||||
sorted_results = sorted(data["results"], key=lambda x: x["index"])
|
||||
return [result["relevance_score"] for result in sorted_results]
|
||||
if 'results' in data:
|
||||
sorted_results = sorted(data['results'], key=lambda x: x['index'])
|
||||
return [result['relevance_score'] for result in sorted_results]
|
||||
else:
|
||||
log.error("No results found in external reranking response")
|
||||
log.error('No results found in external reranking response')
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
log.exception(f"Error in external reranking: {e}")
|
||||
log.exception(f'Error in external reranking: {e}')
|
||||
return None
|
||||
|
||||
Reference in New Issue
Block a user