fahmiaziz98 commited on
Commit
376886a
·
1 Parent(s): d57816a

check linting

Browse files
src/api/routers/embedding.py CHANGED
@@ -25,16 +25,18 @@ from src.core.exceptions import (
25
  ValidationError,
26
  )
27
  from src.api.dependencies import get_model_manager
28
- from src.utils.validators import extract_embedding_kwargs, validate_texts, count_tokens_batch
 
 
 
 
29
  from src.config.settings import get_settings
30
 
31
 
32
  router = APIRouter(tags=["embeddings"])
33
 
34
 
35
- def _ensure_model_type(
36
- config, expected_type: str, model_id: str
37
- ) -> None:
38
  """
39
  Validate that the model configuration matches the expected type.
40
 
@@ -206,4 +208,3 @@ async def create_sparse_embedding(
206
  status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
207
  detail=f"Failed to create query embedding: {str(e)}",
208
  )
209
-
 
25
  ValidationError,
26
  )
27
  from src.api.dependencies import get_model_manager
28
+ from src.utils.validators import (
29
+ extract_embedding_kwargs,
30
+ validate_texts,
31
+ count_tokens_batch,
32
+ )
33
  from src.config.settings import get_settings
34
 
35
 
36
  router = APIRouter(tags=["embeddings"])
37
 
38
 
39
+ def _ensure_model_type(config, expected_type: str, model_id: str) -> None:
 
 
40
  """
41
  Validate that the model configuration matches the expected type.
42
 
 
208
  status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
209
  detail=f"Failed to create query embedding: {str(e)}",
210
  )
 
src/api/routers/rerank.py CHANGED
@@ -6,7 +6,6 @@ It accepts a list of documents and returns a ranked list based on relevance to t
6
  """
7
 
8
  import time
9
- from typing import List
10
  from fastapi import APIRouter, Depends, HTTPException, status
11
  from loguru import logger
12
 
@@ -21,7 +20,7 @@ from src.core.exceptions import (
21
  from src.api.dependencies import get_model_manager
22
  from src.utils.validators import extract_embedding_kwargs
23
 
24
- router = APIRouter(prefix="/rerank",tags=["rerank"])
25
 
26
 
27
  @router.post(
@@ -91,20 +90,16 @@ async def rerank_documents(
91
  processing_time = time.time() - start
92
 
93
  results = []
94
-
95
  for rank_result in ranking_results:
96
- doc_idx = rank_result.get('corpus_id', 0)
97
  if doc_idx < len(valid_docs):
98
  original_idx = valid_docs[doc_idx][0] # Original index
99
  doc_text = documents_list[doc_idx]
100
- score = rank_result['score']
101
-
102
  results.append(
103
- RerankResult(
104
- text=doc_text,
105
- score=score,
106
- index=original_idx
107
- )
108
  )
109
 
110
  logger.info(
@@ -130,4 +125,4 @@ async def rerank_documents(
130
  raise HTTPException(
131
  status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
132
  detail=f"Failed to rerank documents: {str(e)}",
133
- )
 
6
  """
7
 
8
  import time
 
9
  from fastapi import APIRouter, Depends, HTTPException, status
10
  from loguru import logger
11
 
 
20
  from src.api.dependencies import get_model_manager
21
  from src.utils.validators import extract_embedding_kwargs
22
 
23
+ router = APIRouter(prefix="/rerank", tags=["rerank"])
24
 
25
 
26
  @router.post(
 
90
  processing_time = time.time() - start
91
 
92
  results = []
93
+
94
  for rank_result in ranking_results:
95
+ doc_idx = rank_result.get("corpus_id", 0)
96
  if doc_idx < len(valid_docs):
97
  original_idx = valid_docs[doc_idx][0] # Original index
98
  doc_text = documents_list[doc_idx]
99
+ score = rank_result["score"]
100
+
101
  results.append(
102
+ RerankResult(text=doc_text, score=score, index=original_idx)
 
 
 
 
103
  )
104
 
105
  logger.info(
 
125
  raise HTTPException(
126
  status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
127
  detail=f"Failed to rerank documents: {str(e)}",
128
+ )
src/core/base.py CHANGED
@@ -6,7 +6,7 @@ must follow, ensuring consistency across dense and sparse embeddings.
6
  """
7
 
8
  from abc import ABC, abstractmethod
9
- from typing import Any, Dict, List, Optional, Union
10
 
11
 
12
  class BaseEmbeddingModel(ABC):
 
6
  """
7
 
8
  from abc import ABC, abstractmethod
9
+ from typing import Any, Dict, List, Union
10
 
11
 
12
  class BaseEmbeddingModel(ABC):
src/models/embeddings/dense.py CHANGED
@@ -112,7 +112,7 @@ class DenseEmbeddingModel(BaseEmbeddingModel):
112
 
113
  try:
114
  embeddings = self.model.encode(input, **kwargs)
115
-
116
  return [
117
  emb.tolist() if hasattr(emb, "tolist") else list(emb)
118
  for emb in embeddings
 
112
 
113
  try:
114
  embeddings = self.model.encode(input, **kwargs)
115
+
116
  return [
117
  emb.tolist() if hasattr(emb, "tolist") else list(emb)
118
  for emb in embeddings
src/models/embeddings/rank.py CHANGED
@@ -113,22 +113,17 @@ class RerankModel:
113
  """
114
  if not self._loaded or self.model is None:
115
  self.load()
116
-
117
  try:
118
- ranking_results = self.model.rank(
119
- query,
120
- documents,
121
- top_k=top_k,
122
- **kwargs
123
- )
124
-
125
  # Normalize scores to 0-1 range for consistency
126
  normalized_results = self._normalize_rerank_scores(ranking_results)
127
-
128
  logger.debug(
129
  f"Reranked {len(documents)} docs, returned top {len(normalized_results)}"
130
  )
131
-
132
  return normalized_results
133
 
134
  except Exception as e:
@@ -137,9 +132,7 @@ class RerankModel:
137
  raise RerankingDocumentError(self.model_id, error_msg)
138
 
139
  def _normalize_rerank_scores(
140
- self,
141
- rankings: List[Dict],
142
- target_range: tuple = (0, 1)
143
  ) -> List[Dict]:
144
  """
145
  Normalize reranking scores using min-max normalization.
@@ -154,35 +147,30 @@ class RerankModel:
154
  """
155
  if not rankings:
156
  return []
157
-
158
  raw_scores = [ranking["score"] for ranking in rankings]
159
-
160
  min_score = min(raw_scores)
161
  max_score = max(raw_scores)
162
-
163
  if max_score == min_score:
164
  return [
165
- {
166
- "corpus_id": r["corpus_id"],
167
- "score": target_range[1]
168
- }
169
  for r in rankings
170
  ]
171
-
172
  target_min, target_max = target_range
173
  normalized_rankings = []
174
-
175
  for ranking in rankings:
176
  score = ranking["score"]
177
- normalized_score = (
178
- target_min +
179
- (score - min_score) * (target_max - target_min) / (max_score - min_score)
 
 
180
  )
181
- normalized_rankings.append({
182
- "corpus_id": ranking["corpus_id"],
183
- "score": float(normalized_score)
184
- })
185
-
186
  return normalized_rankings
187
 
188
  @property
@@ -222,4 +210,4 @@ class RerankModel:
222
  f"id={self.model_id}, "
223
  f"type={self.model_type}, "
224
  f"loaded={self.is_loaded})"
225
- )
 
113
  """
114
  if not self._loaded or self.model is None:
115
  self.load()
116
+
117
  try:
118
+ ranking_results = self.model.rank(query, documents, top_k=top_k, **kwargs)
119
+
 
 
 
 
 
120
  # Normalize scores to 0-1 range for consistency
121
  normalized_results = self._normalize_rerank_scores(ranking_results)
122
+
123
  logger.debug(
124
  f"Reranked {len(documents)} docs, returned top {len(normalized_results)}"
125
  )
126
+
127
  return normalized_results
128
 
129
  except Exception as e:
 
132
  raise RerankingDocumentError(self.model_id, error_msg)
133
 
134
  def _normalize_rerank_scores(
135
+ self, rankings: List[Dict], target_range: tuple = (0, 1)
 
 
136
  ) -> List[Dict]:
137
  """
138
  Normalize reranking scores using min-max normalization.
 
147
  """
148
  if not rankings:
149
  return []
150
+
151
  raw_scores = [ranking["score"] for ranking in rankings]
152
+
153
  min_score = min(raw_scores)
154
  max_score = max(raw_scores)
155
+
156
  if max_score == min_score:
157
  return [
158
+ {"corpus_id": r["corpus_id"], "score": target_range[1]}
 
 
 
159
  for r in rankings
160
  ]
161
+
162
  target_min, target_max = target_range
163
  normalized_rankings = []
164
+
165
  for ranking in rankings:
166
  score = ranking["score"]
167
+ normalized_score = target_min + (score - min_score) * (
168
+ target_max - target_min
169
+ ) / (max_score - min_score)
170
+ normalized_rankings.append(
171
+ {"corpus_id": ranking["corpus_id"], "score": float(normalized_score)}
172
  )
173
+
 
 
 
 
174
  return normalized_rankings
175
 
176
  @property
 
210
  f"id={self.model_id}, "
211
  f"type={self.model_type}, "
212
  f"loaded={self.is_loaded})"
213
+ )
src/models/schemas/requests.py CHANGED
@@ -13,7 +13,7 @@ from .common import EmbeddingOptions
13
  class BaseEmbedRequest(BaseModel):
14
  """
15
  OpenAI-compatible embedding request.
16
-
17
  Matches the format of OpenAI's embeddings API:
18
  https://platform.openai.com/docs/api-reference/embeddings
19
  """
@@ -25,16 +25,11 @@ class BaseEmbedRequest(BaseModel):
25
  )
26
 
27
  encoding_format: Optional[Literal["float", "base64"]] = Field(
28
- default="float",
29
- description="Encoding format"
30
  )
31
- dimensions: Optional[int] = Field(
32
- None,
33
- description="Output dimensions")
34
 
35
- user: Optional[str] = Field(
36
- None,
37
- description="User identifier")
38
 
39
  options: Optional[EmbeddingOptions] = Field(
40
  None, description="Optional embedding generation parameters"
 
13
  class BaseEmbedRequest(BaseModel):
14
  """
15
  OpenAI-compatible embedding request.
16
+
17
  Matches the format of OpenAI's embeddings API:
18
  https://platform.openai.com/docs/api-reference/embeddings
19
  """
 
25
  )
26
 
27
  encoding_format: Optional[Literal["float", "base64"]] = Field(
28
+ default="float", description="Encoding format"
 
29
  )
30
+ dimensions: Optional[int] = Field(None, description="Output dimensions")
 
 
31
 
32
+ user: Optional[str] = Field(None, description="User identifier")
 
 
33
 
34
  options: Optional[EmbeddingOptions] = Field(
35
  None, description="Optional embedding generation parameters"
src/models/schemas/responses.py CHANGED
@@ -23,6 +23,7 @@ class BaseEmbedResponse(BaseModel):
23
 
24
  class EmbeddingObject(BaseModel):
25
  """Single embedding object."""
 
26
  object: Literal["embedding"] = "embedding"
27
  embedding: List[float] = Field(..., description="Embedding vector")
28
  index: int = Field(..., description="Index of the embedding")
@@ -30,6 +31,7 @@ class EmbeddingObject(BaseModel):
30
 
31
  class TokenUsage(BaseModel):
32
  """Usage statistics."""
 
33
  prompt_tokens: int
34
  total_tokens: int
35
 
@@ -44,14 +46,15 @@ class DenseEmbedResponse(BaseEmbedResponse):
44
  data: List of generated dense embeddings
45
  model: Identifier of the model used
46
  usage: Usage statistics
47
-
48
  """
 
49
  object: Literal["list"] = "list"
50
  data: List[EmbeddingObject]
51
  model: str = Field(..., description="Model identifier used")
52
  usage: TokenUsage = Field(..., description="Usage statistics")
53
 
54
- class Config:
55
  json_schema_extra = {
56
  "example": {
57
  "object": "list",
 
23
 
24
  class EmbeddingObject(BaseModel):
25
  """Single embedding object."""
26
+
27
  object: Literal["embedding"] = "embedding"
28
  embedding: List[float] = Field(..., description="Embedding vector")
29
  index: int = Field(..., description="Index of the embedding")
 
31
 
32
  class TokenUsage(BaseModel):
33
  """Usage statistics."""
34
+
35
  prompt_tokens: int
36
  total_tokens: int
37
 
 
46
  data: List of generated dense embeddings
47
  model: Identifier of the model used
48
  usage: Usage statistics
49
+
50
  """
51
+
52
  object: Literal["list"] = "list"
53
  data: List[EmbeddingObject]
54
  model: str = Field(..., description="Model identifier used")
55
  usage: TokenUsage = Field(..., description="Usage statistics")
56
 
57
+ class Config:
58
  json_schema_extra = {
59
  "example": {
60
  "object": "list",