Justin44 commited on
Commit
e7737d0
·
verified ·
1 Parent(s): 05f672b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +10 -7
app.py CHANGED
@@ -1,3 +1,6 @@
 
 
 
1
  import torch
2
  from fastapi import FastAPI
3
  from pydantic import BaseModel
@@ -32,7 +35,7 @@ splade_model = None
32
  app = FastAPI(
33
  title="Hybrid Vector Generation API",
34
  description="An API to generate dense and sparse vectors for a given text query.",
35
- version="1.0.0"
36
  )
37
 
38
  # --- Pydantic Models for API ---
@@ -46,7 +49,7 @@ class SparseVectorResponse(BaseModel):
46
  values: list[float]
47
 
48
  class VectorResponse(BaseModel):
49
- """The final JSON response model containing both vectors."""
50
  dense_vector: list[float]
51
  sparse_vector: SparseVectorResponse
52
 
@@ -64,10 +67,10 @@ async def load_models():
64
  dense_model = SentenceTransformer(DENSE_MODEL_ID, device=DEVICE)
65
  splade_tokenizer = AutoTokenizer.from_pretrained(SPLADE_QUERY_MODEL_ID)
66
  splade_model = AutoModelForMaskedLM.from_pretrained(SPLADE_QUERY_MODEL_ID).to(DEVICE)
67
- logger.info("Great news! All models have been loaded successfully.")
68
  except Exception as e:
69
  logger.critical(f"Oh no, a critical error occurred while loading models: {e}", exc_info=True)
70
- # In a real-world scenario, you might want the app to fail startup if models don't load.
71
  raise e
72
 
73
  def compute_splade_vector(text: str) -> models.SparseVector:
@@ -119,7 +122,7 @@ async def vectorize_query(request: QueryRequest):
119
  logger.info(f"The incoming search query from n8n is: '{request.query_text}'")
120
 
121
  # 1. Generate Dense Vector
122
- logger.info("First, I'm generating the dense vector for semantic meaning...")
123
  dense_query_vector = dense_model.encode(request.query_text).tolist()
124
  logger.info("Done with the dense vector. It has %d dimensions.", len(dense_query_vector))
125
  logger.info("Here's a small sample of the dense vector: %s...", str(dense_query_vector[:4]))
@@ -131,8 +134,8 @@ async def vectorize_query(request: QueryRequest):
131
  logger.info("Here's a sample of the sparse vector indices: %s...", str(sparse_query_vector.indices[:4]))
132
 
133
  # 3. Construct and return the response
134
- logger.info("Everything looks good. I'm packaging up the vectors to send back.")
135
- logger.info("=========================================================")
136
 
137
  final_response = VectorResponse(
138
  dense_vector=dense_query_vector,
 
1
+ # Author : Justin
2
+ # Program : Vectorizer for Hybrid Search
3
+ # Instructions : Check README.md
4
  import torch
5
  from fastapi import FastAPI
6
  from pydantic import BaseModel
 
35
  app = FastAPI(
36
  title="Hybrid Vector Generation API",
37
  description="An API to generate dense and sparse vectors for a given text query.",
38
+ version="1.2.0"
39
  )
40
 
41
  # --- Pydantic Models for API ---
 
49
  values: list[float]
50
 
51
  class VectorResponse(BaseModel):
52
+ """Final JSON response model containing both vectors."""
53
  dense_vector: list[float]
54
  sparse_vector: SparseVectorResponse
55
 
 
67
  dense_model = SentenceTransformer(DENSE_MODEL_ID, device=DEVICE)
68
  splade_tokenizer = AutoTokenizer.from_pretrained(SPLADE_QUERY_MODEL_ID)
69
  splade_model = AutoModelForMaskedLM.from_pretrained(SPLADE_QUERY_MODEL_ID).to(DEVICE)
70
+ logger.info("YAaay! All models have been loaded successfully.")
71
  except Exception as e:
72
  logger.critical(f"Oh no, a critical error occurred while loading models: {e}", exc_info=True)
73
+ # In a real-world scenario, you might want the app to fail startup if models don't load
74
  raise e
75
 
76
  def compute_splade_vector(text: str) -> models.SparseVector:
 
122
  logger.info(f"The incoming search query from n8n is: '{request.query_text}'")
123
 
124
  # 1. Generate Dense Vector
125
+ logger.info("First, generating the dense vector for semantic meaning...")
126
  dense_query_vector = dense_model.encode(request.query_text).tolist()
127
  logger.info("Done with the dense vector. It has %d dimensions.", len(dense_query_vector))
128
  logger.info("Here's a small sample of the dense vector: %s...", str(dense_query_vector[:4]))
 
134
  logger.info("Here's a sample of the sparse vector indices: %s...", str(sparse_query_vector.indices[:4]))
135
 
136
  # 3. Construct and return the response
137
+ logger.info("Everything looks good. Packaging up the vectors to send back.")
138
+ logger.info("-----------------------------------------------------------------")
139
 
140
  final_response = VectorResponse(
141
  dense_vector=dense_query_vector,