Upload 6 files

README.md

@@ -1,10 +1,13 @@
----
-title: Fastapi Ai Endpoints
-emoji: 🐢
-colorFrom: green
-colorTo: red
-sdk: docker
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+https://huggingface.co/blog/HemanthSai7/deploy-applications-on-huggingface-spaces
+
+Objective: Convert any Huggingface repository into an API endpoint
+
+Users should be able to call the task and get back in the standard format
+
+/sentence-embeddings
+
+{
+    "model": "BAAI/bge-base-en-v1.5",
+    "inputs: ["This is one text", "This is second text"],
+    "parameters": {}
+}

app.py

@@ -0,0 +1,20 @@
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi import FastAPI
+import sentence_embeddings
+
+app = FastAPI()
+
+# CORS Support: https://stackoverflow.com/a/66460861
+origins = ["*"]
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+app.include_router(sentence_embeddings.router)
+if __name__ == '__main__':
+    import uvicorn
+    uvicorn.run(app, host='0.0.0.0', port=8000)

config.py

@@ -0,0 +1,6 @@
+import os
+import dotenv
+
+dotenv.load_dotenv()
+
+TEST_MODE = (os.getenv('TEST_MODE', 'False') == "True")

logger.py

@@ -0,0 +1,4 @@
+from datetime import datetime
+
+def log(data: dict):
+    print(f"{datetime.now().isoformat()}: {data}")

requirements.txt

@@ -0,0 +1,4 @@
+transformers
+torch
+fastapi
+uvicorn

sentence_embeddings.py

@@ -0,0 +1,84 @@
+from typing import Optional
+from fastapi import APIRouter
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModel
+import torch
+from datetime import datetime
+from logger import log
+from hf_to_api.config import TEST_MODE
+
+router = APIRouter()
+
+class SentenceEmbeddingsInput(BaseModel):
+    inputs: list[str]
+    model: str
+    parameters: dict
+
+class SentenceEmbeddingsOutput(BaseModel):
+    embeddings: Optional[list[list[float]]] = None
+    error: Optional[str] = None
+
+@router.post('/sentence-embeddings')
+def sentence_embeddings(inputs: SentenceEmbeddingsInput):
+    start_time = datetime.now()
+    fn = sentence_embeddings_mapping.get(inputs.model)
+    if not fn:
+        return SentenceEmbeddingsOutput(
+            error=f'No sentence embeddings model found for {inputs.model}'
+        )
+    
+    try:
+        embeddings = fn(inputs.inputs, inputs.parameters)
+        
+        log({
+            "task": "sentence_embeddings",
+            "model": inputs.model,
+            "start_time": start_time.isoformat(),
+            "time_taken": (datetime.now() - start_time).total_seconds(),
+            "inputs": inputs.inputs,
+            "outputs": embeddings,
+            "parameters": inputs.parameters,
+        })
+        loaded_models_last_updated[inputs.model] = datetime.now()
+        return SentenceEmbeddingsOutput(
+            embeddings=embeddings
+        )
+    except Exception as e:
+        return SentenceEmbeddingsOutput(
+            error=str(e)
+        )
+
+def generic_sentence_embeddings(model_name: str):
+    global loaded_models
+
+    def process_texts(texts: list[str], parameters: dict):
+        if TEST_MODE:
+            return [[0.1,0.2]] * len(texts)
+    
+        if model_name in loaded_models:
+            tokenizer, model = loaded_models[model_name]
+        else:
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
+            model = AutoModel.from_pretrained(model_name)
+            loaded_models[model] = (tokenizer, model)
+
+        # Tokenize sentences
+        encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
+        with torch.no_grad():
+            model_output = model(**encoded_input)
+            sentence_embeddings = model_output[0][:, 0]
+
+        # normalize embeddings
+        sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
+        return sentence_embeddings.tolist()
+
+    return process_texts
+
+# Polling every X minutes to 
+loaded_models = {}
+loaded_models_last_updated = {}
+
+sentence_embeddings_mapping = {
+    'BAAI/bge-base-en-v1.5': generic_sentence_embeddings('BAAI/bge-base-en-v1.5'),
+    'BAAI/bge-large-en-v1.5': generic_sentence_embeddings('BAAI/bge-large-en-v1.5'),
+}