diff --git a/main.py b/main.py index 649b826..ff5e4ac 100644 --- a/main.py +++ b/main.py @@ -20,7 +20,7 @@ from engine import FS app = FastAPI() data = pd.DataFrame() - +mlb = MultiLabelBinarizer() def inference(first: pandas.core.series.Series, second_id: str, df=None): if df is not None: @@ -59,14 +59,41 @@ def process_dataframe(df, production): @app.on_event('startup') async def startup_event(): global data + global mlb data = pd.read_csv('processed_data.csv', index_col='id', converters={'genres': pd.eval}) all_genres = data.genres.explode().unique() - mlb = MultiLabelBinarizer() mlb.fit([all_genres]) data['genres'] = data['genres'].apply(lambda x: mlb.transform([x])[0]) data['emotions'] = data[['Happy', 'Angry', 'Surprise', 'Sad', 'Fear']].values.tolist() +@app.get('/find/{title}') +def titles(title: str): + response = {} + for index, row in data.iterrows(): + if title.lower() in row['title'].lower(): + response[index] = {'title': row['title'], 'year': row['release_year']} + return response + + +@app.get('/details/{production_id}') +def details(production_id: str): + try: + production = data.loc[production_id] + except: + return {'error': f'{production_id} is not a valid id'} + genres = production['genres'] + genres = mlb.inverse_transform(genres.reshape(1, -1))[0] + return { + 'title': production['title'], + 'type': production['type'], + 'description': production['description'], + 'year': int(production['release_year']), + 'runtime': int(production['runtime']), + 'genres': genres, + } + + @app.get('/score/{first_id}/{second_id}') def rec_score(first_id: str, second_id: str): try: