Update app.py

app.py

@@ -2,19 +2,15 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from sentence_transformers import SentenceTransformer, util
 from transformers import pipeline
-#from transformers import T5Tokenizer, T5ForConditionalGeneration
+import numpy as np
 
 
-# Initialize FastAPI app
 app = FastAPI()
 
-# Load models
 model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 question_model = "deepset/tinyroberta-squad2"
 nlp = pipeline('question-answering', model=question_model, tokenizer=question_model)
 
-#t5tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
-#t5model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large")
 summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
 
 # Define request models
@@ -80,6 +76,31 @@ async def t5answer(request: T5QuestionRequest):
     resp = summarizer(request.context, max_length=130, min_length=30, do_sample=False)
     return T5Response(answer = resp[0]["summary_text"])
 
+
+# Define API endpoints
+@app.post("/modify_query2", response_model=ModifyQueryResponse)
+async def modify_query(request: ModifyQueryRequest):
+    try:
+        embeddings = optimize_embedding([request.query_string])
+        return ModifyQueryResponse(embeddings=embeddings[0].tolist())
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+def optimize_embedding(texts, precision='uint8'):
+    # Step 1: Generate embeddings with 384 dimensions
+    embeddings = model.encode(texts)
+
+    # Step 2: Quantize embeddings to chosen precision (e.g., uint8)
+    if precision == 'uint8':
+        quantized_embeddings = np.array(embeddings, dtype='float32').astype('uint8')
+    elif precision == 'uint16':
+        quantized_embeddings = np.array(embeddings, dtype='float32').astype('uint16')
+    else:
+        raise ValueError("Unsupported precision. Use 'uint8' or 'uint16'.")
+
+    return quantized_embeddings
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)