Den-d3j2d commited on
Commit
8fed2d7
·
verified ·
1 Parent(s): 542680b

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +23 -0
app.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI
2
+ from pydantic import BaseModel
3
+ from typing import List
4
+ import torch
5
+ from transformers import AutoTokenizer, AutoModel
6
+ from src.model.encoder import ProdFeatureEncoder
7
+ from src.config import ModelConfig
8
+
9
+ app = FastAPI()
10
+
11
+ class TextInput(BaseModel):
12
+ text: str
13
+
14
+ class EmbeddingOutput(BaseModel):
15
+ embedding: List[float]
16
+
17
+ @app.post("/encode_text", response_model=EmbeddingOutput)
18
+ async def encode_text(text_input: TextInput):
19
+ config = ModelConfig()
20
+ model = ProdFeatureEncoder(config=config)
21
+ embedding = model(text_input.text)
22
+ return {"embedding": embedding.tolist()}
23
+