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

# Initialize the model
model = SentenceTransformer("PartAI/Tooka-SBERT")

# Create the FastAPI app
app = FastAPI()

# Pydantic model for input data
class TextInput(BaseModel):
    sentences: List[str]

@app.get('/')
def index():
    return {'message': 'Sentence embedding API.'}

# Endpoint to get embeddings
@app.post("/get_embeddings")
async def get_embeddings(input_data: TextInput):
    # Get embeddings for the input sentences
    embeddings = model.encode(input_data.sentences)
    return {"embeddings": embeddings.tolist()}

# To run the app, save this code to a file, and then run `uvicorn filename:app --reload`