File size: 1,104 Bytes
b48bc49 |
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 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
from fastapi import FastAPI
from pydantic import BaseModel
from utils import compare_embeddings
from models.esm_2_650m import get_embedding as get_embedding_esm_2_650m
app = FastAPI()
class CompareRequest(BaseModel):
sequence_1: str
sequence_2: str
model: str = "esm_2_650m"
model_mapping = {"esm_2_650m": get_embedding_esm_2_650m}
# ----------------------------------------------------------------------
@app.get("/")
def root():
return {
"message": "API is running. Use POST /compare to compare protein sequences."
}
# ----------------------------------------------------------------------
@app.post("/compare")
def compare(request: CompareRequest):
model = request.model
if model not in model_mapping:
return {"error": "Model not supported"}
emb1 = model_mapping[model](request.sequence_1)
emb2 = model_mapping[model](request.sequence_2)
similarity, classification = compare_embeddings(emb1, emb2)
return {
"cosine_similarity": float(similarity),
"classification": classification,
"model": model,
}
|