itod/embeddings

1import os2from fastapi import FastAPI3from typing import Union, List, Dict, Tuple, Optional4from pydantic import BaseModel, Field5from angle_emb import AnglE6 7class EmbeddingInput(BaseModel):8    input: Union[List[str], Tuple[str], List[Dict], str] = Field(..., description="The input to be encoded")9    model: Optional[str] = None10    encoding_format: Optional[str] = 'float'11    dimensions: Optional[int] = None12    user: Optional[str] = None13 14app = FastAPI()15 16# Get the model name and path from the environment variables17model_name = os.getenv('MODEL_NAME', default='WhereIsAI/UAE-Large-V1')18model_path = os.getenv('MODEL_PATH', default='models/WhereIsAI/UAE-Large-V1')19 20# Load the model21try:22    angle_model = AnglE.from_pretrained(model_path, pooling_strategy='cls').to('cpu')23except Exception as e:24    print(f"Failed to load model from path {model_path}. Error: {str(e)}")25 26@app.get("/")27def read_root():28    return {29        "model_name": model_name,30        "model_path": model_path,31        "message": "Model is up and running",32        "route_info": {33            "/": "Returns the model info",34            "/health": "Returns the health status of the application",35            "/v1/embeddings": 'POST route to get embeddings. Usage: curl -H "Content-Type: application/json" -d \'{ "input": "Your text string goes here" }\' http://localhost:8080/v1/embeddings'36        }37    }38 39@app.get("/health")40def health_check():41    return {"health": "ok"}42 43@app.post("/v1/embeddings")44def get_embeddings(embedding_input: EmbeddingInput):45    # # Check if the input is an empty string46    # if not embedding_input.input.strip():47    #     return {48    #         "object": "list",49    #         "data": [],50    #         "model": model_name,51    #         "usage": {"prompt_tokens": 0, "total_tokens": 0},52    #     }53 54    # Encode the input text using the model55    embeddings = angle_model.encode(embedding_input.input, embedding_size=embedding_input.dimensions)56 57    # Create a response format compatible with OpenAI's API58    response = {59        "object": "list",60        "data": [61            {"object": "embedding", "index": i, "embedding": emb.tolist()}62            for i, emb in enumerate(embeddings)63        ],64        "model": model_name,65        "usage": {"prompt_tokens": len(embedding_input.input), "total_tokens": len(embedding_input.input)},66    }67 68    return response69 70if __name__ == "__main__":71    import uvicorn72    uvicorn.run(app, host="0.0.0.0", port=8080)73