cosine

main.py

@@ -1,34 +1,72 @@
-from fastapi import FastAPI
+rom fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from sentence_transformers import SentenceTransformer
+from sentence_transformers import SentenceTransformer, util
 import numpy as np
 
-import os
-os.environ["HF_HOME"] = "/tmp/huggingface"
-# Initialize FastAPI app
+# Initialize the FastAPI app
 app = FastAPI()
 
-from sentence_transformers import SentenceTransformer
+# Load the embedding models
+embedding_model = SentenceTransformer("Alibaba-NLP/gte-base-en-v1.5", trust_remote_code=True)
+similarity_model = SentenceTransformer("all-MiniLM-L6-v2")
 
-embedding_model = SentenceTransformer("Alibaba-NLP/gte-base-en-v1.5")
+# Define request body schemas
+class TextInput(BaseModel):
+    text: str
 
+class SimilarityInput(BaseModel):
+    text1: str
+    text2: str
 
-# Define request body structure
-class TextRequest(BaseModel):
-    text: str
+# Home route
+@app.get("/")
+async def home():
+    return {"message": "Welcome to the embedding and similarity API. Use /docs to test the endpoints."}
+
+# Endpoint for generating embeddings
+@app.post("/embed")
+async def generate_embedding(text_input: TextInput):
+    """
+    Generate a 768-dimensional embedding for the input text.
+    Returns the embedding in a structured format with rounded values.
+    """
+    try:
+        embedding = embedding_model.encode(text_input.text, convert_to_tensor=True).cpu().numpy()
+        rounded_embedding = np.round(embedding, decimals=2).tolist()
+        dimensions = len(rounded_embedding)
+
+        return {
+            "dimensions": dimensions,
+            "embeddings": rounded_embedding
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+# New endpoint for calculating cosine similarity
+@app.post("/similarity")
+async def calculate_similarity(similarity_input: SimilarityInput):
+    """
+    Calculate cosine similarity between two text inputs.
+    """
+    try:
+        # Compute embeddings
+        embeddings1 = similarity_model.encode(similarity_input.text1, convert_to_tensor=True)
+        embeddings2 = similarity_model.encode(similarity_input.text2, convert_to_tensor=True)
 
-# Define response structure
-class EmbeddingResponse(BaseModel):
-    dimensions: int
-    embedding: list[float]
+        # Compute cosine similarity
+        cosine_similarity = util.cos_sim(embeddings1, embeddings2).item()
 
-# Create API endpoint
-@app.post("/get_embedding", response_model=EmbeddingResponse)
-async def get_embedding(request: TextRequest):
-    # Generate embedding
-    embedding = embedding_model.encode([request.text])[0]  # Extract first item
+        return {
+            "text1": similarity_input.text1,
+            "text2": similarity_input.text2,
+            "cosine_similarity": round(cosine_similarity, 4)
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
-    # Convert to list and return response
-    return {"dimensions": len(embedding), "embedding": embedding.tolist()}
+# Run the FastAPI app
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)