Spaces:
Running
Running
Update app.py
Browse files
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.
|
36 |
)
|
37 |
|
38 |
# --- Pydantic Models for API ---
|
@@ -46,7 +49,7 @@ class SparseVectorResponse(BaseModel):
|
|
46 |
values: list[float]
|
47 |
|
48 |
class VectorResponse(BaseModel):
|
49 |
-
"""
|
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("
|
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,
|
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.
|
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,
|