|
from fastapi import FastAPI |
|
from sentence_transformers import SentenceTransformer |
|
|
|
|
|
app = FastAPI( |
|
title="Generate Embeddings API", |
|
description=""" |
|
## Generate Embeddings API |
|
A simple API to generate sentence embeddings using **Sentence Transformers**. |
|
|
|
### π Useful Links |
|
- [Hugging Face Code](https://huggingface.co/spaces/TheMihirNaik/generate-embeddings-api/tree/main) |
|
- [SEO Workflow Automation Newsletter](https://mihirnaik.substack.com/) |
|
- [Mihir Naik](https://www.mihirnaik.com) |
|
|
|
### How It Works |
|
1. Send a **POST request** to `/generate-embeddings/` with a list of sentences. |
|
2. Get **embeddings** as JSON response. |
|
|
|
Built by **Mihir Naik** π |
|
""", |
|
version="1.0.0", |
|
contact={ |
|
"name": "Mihir Naik", |
|
"url": "https://www.mihirnaik.com", |
|
} |
|
) |
|
|
|
|
|
all_MiniLM_L6_V2_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2') |
|
|
|
intfloat_e5_large_v2_model = SentenceTransformer('intfloat/e5-large-v2') |
|
|
|
@app.get("/") |
|
def redirect_to_docs(): |
|
""" |
|
Redirects to the FastAPI documentation. |
|
|
|
Returns: |
|
RedirectResponse: A response that redirects to the documentation URL. |
|
""" |
|
from fastapi.responses import RedirectResponse |
|
return RedirectResponse(url="/docs") |
|
|
|
from typing import List |
|
|
|
@app.post("/generate-embeddings/sentence-transformers-all-MiniLM-L6-v2/") |
|
def generate_embeddings_all_MiniLM_L6_V2_model(sentences: List[str]): |
|
""" |
|
Generates embeddings for a list of sentences using the MiniLM model. |
|
|
|
Args: |
|
sentences (List[str]): A list of sentences to generate embeddings for. |
|
|
|
Returns: |
|
dict: A dictionary containing the embeddings as a JSON-compatible list. |
|
""" |
|
embeddings = all_MiniLM_L6_V2_model.encode(sentences) |
|
return {"embeddings": embeddings.tolist()} |
|
|
|
|
|
@app.post("/generate-embeddings/intfloat-e5-large-v2/") |
|
def generate_embeddings_intfloat_e5_large_v2_model(sentences: List[str]): |
|
""" |
|
Generates embeddings for a list of sentences using the intfloat/e5-large-v2 model. |
|
|
|
Args: |
|
sentences (List[str]): A list of sentences to generate embeddings for. |
|
|
|
Returns: |
|
dict: A dictionary containing the embeddings as a JSON-compatible list. |
|
""" |
|
embeddings = intfloat_e5_large_v2_model.encode(sentences) |
|
return {"embeddings": embeddings.tolist()} |
|
|