from dataset_util import load_data, get_num_rows
from fastapi import FastAPI

twitter_text_dataset = load_data("bittensor-dataset/twitter-text-dataset")
twitter_text_num_rows = get_num_rows(twitter_text_dataset)
twitter_image_dataset = load_data("bittensor-dataset/twitter-image-dataset")
twitter_image_num_rows = get_num_rows(twitter_image_dataset)
app = FastAPI()

@app.get("/")
def read_root():
    return {"twitter_text_rows": twitter_text_num_rows, "twitter_image_rows": twitter_image_num_rows}