File size: 2,007 Bytes
ceaeaf3 |
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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
# api/main.py
from fastapi import FastAPI, HTTPException
import sentence_transformers
from huggingface_hub import hf_hub_download
import pandas as pd
from src.processor import send_to_dataset,search_and_retrieve,generate_tech
from typing import List, Dict
from pydantic import BaseModel
from datasets import load_dataset
# This is the main application object that Uvicorn will run
app = FastAPI(
title="My Standalone API",
description="An API hosted on Hugging Face Spaces",
version="1.0.0"
)
model = sentence_transformers.SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
dataset = load_dataset("heymenn/Technologies", split="train")
class SearchInput(BaseModel):
title: str
class SearchOutput(BaseModel):
title: str
purpose: str
score: float
top5: List[Dict]
class GenerateInput(BaseModel):
title: str
instructions: str
force: bool = False
class GenerateOutput(BaseModel):
name: str
purpose: str
problem_types_solved: str
advantages: str
limitations: str
domain_tags: str
@app.post("/search-technologies", response_model=SearchOutput)
def post_search(payload: SearchInput):
"""
Endpoint that returns a search result.
"""
config = {"dataset": dataset, "model": model}
res = search_and_retrieve(payload.title, config)
return res
@app.post("/generate-technology", response_model=GenerateOutput)
def post_generate_and_push(payload: GenerateInput):
"""
Endpoint to generate a technology and push it to the dataset
"""
config = {"dataset": dataset, "model": model}
res = search_and_retrieve(payload.title, config)
if res["score"] >= 0.7 and not payload.force:
raise HTTPException(status_code=500, detail=f"Cannot generate the technology a high score of {res['score']} have been found for the technology : {res['title']}")
json_response = generate_tech(payload.title, payload.instructions)
send_to_dataset(json_response, model)
return json_response |