from fastapi import FastAPI
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer

import os
os.environ["HF_HOME"] = "/tmp/huggingface"

# Initialize FastAPI app
app = FastAPI()

# Load pretrained model from Hugging Face (instead of hf_hub_download)
model = SentenceTransformer("all-MiniLM-L6-v2")  # Updated model

# Define request structure
class TextRequest(BaseModel):
    text: str

# Define response structure
class EmbeddingResponse(BaseModel):
    dimensions: int
    embedding: list[float]

# Endpoint to get text embedding
@app.post("/get_embedding", response_model=EmbeddingResponse)
async def get_embedding(request: TextRequest):
    embedding = model.encode([request.text])[0]  # Generate embedding
    return {"dimensions": len(embedding), "embedding": embedding.tolist()}