Spaces:
Sleeping
Sleeping
File size: 598 Bytes
8fed2d7 ac50ec3 8fed2d7 f1ddd83 6bdf3ab f1ddd83 8fed2d7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 |
from fastapi import FastAPI
from pydantic import BaseModel
from typing import List
import torch
from transformers import AutoTokenizer, AutoModel
from src.model.encoder import ProdFeatureEncoder
from src.config.config import ModelConfig
app = FastAPI()
class TextInput(BaseModel):
text: str
class EmbeddingOutput(BaseModel):
embedding: List[float]
config = ModelConfig()
model = ProdFeatureEncoder(config=config)
@app.get("/encode_text/{text}", response_model=EmbeddingOutput)
async def encode_text(text: str):
embedding = model(text)
return {"embedding": embedding.tolist()}
|