forked from quickcatch/CourseProject
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathapi.py
More file actions
30 lines (25 loc) · 878 Bytes
/
api.py
File metadata and controls
30 lines (25 loc) · 878 Bytes
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
from sklearn.cluster import KMeans
from fastapi import FastAPI
from clustering import classify, CosSimilarity, get_similar_docs
from joblib import load
from sklearn.feature_extraction.text import TfidfVectorizer
import os
app = FastAPI()
cluster_dir = "clusters"
data_dir = "data_entire"
url_df = None
model = None
vectorizer = None
cos = None
@app.on_event("startup")
async def startup():
global model, url_df, cos, vectorizer
model = load(os.path.join(cluster_dir, "model.pkl"))
url_df = load(os.path.join(cluster_dir, "dataframe.pkl"))
vectorizer = load(os.path.join(cluster_dir, "vectorizer.pkl"))
cos = CosSimilarity(data_dir, model, url_df, vectorizer)
@app.get("/classify/{url:path}")
async def read_item(url):
global cos
docs, titles = get_similar_docs(url, cos)
return {"docs": docs, "titles": titles}