!pip install -q aiohttp==3.8.3 datasets==2.14.6 pandas==1.5.3 requests==2.31.0 tqdm==4.66.1 huggingface-hub>=0.20

import asyncio
from getpass import getpass
import json
from pathlib import Path
import time
from typing import Optional

from aiohttp import ClientSession, ClientTimeout
from datasets import load_dataset, Dataset, DatasetDict
from huggingface_hub import notebook_login, create_inference_endpoint, list_inference_endpoints, whoami
import numpy as np
import pandas as pd
import requests
from tqdm.auto import tqdm

DATASET_IN = "derek-thomas/dataset-creator-reddit-bestofredditorupdates"
DATASET_OUT = "processed-subset-bestofredditorupdates"
ENDPOINT_NAME = "boru-jina-embeddings-demo-ie"

MAX_WORKERS = 5  # This is for how many async workers you want. Choose based on the model and hardware
ROW_COUNT = 100  # Choose None to use all rows, Im using 100 just for a demo

# GPU Choice
VENDOR = "aws"
REGION = "us-east-1"
INSTANCE_SIZE = "medium"
INSTANCE_TYPE = "g5.2xlarge"

notebook_login()

>>> who = whoami()
>>> organization = getpass(
...     prompt="What is your Hugging Face 🤗 username or organization? (with an added payment method)"
... )

>>> namespace = organization or who["name"]

dataset = load_dataset(DATASET_IN)
dataset["train"]

documents = dataset["train"].to_pandas().to_dict("records")[:ROW_COUNT]
len(documents), documents[0]

try:
    endpoint = create_inference_endpoint(
        ENDPOINT_NAME,
        repository="jinaai/jina-embeddings-v2-base-en",
        revision="7302ac470bed880590f9344bfeee32ff8722d0e5",
        task="sentence-embeddings",
        framework="pytorch",
        accelerator="gpu",
        instance_size=INSTANCE_SIZE,
        instance_type=INSTANCE_TYPE,
        region=REGION,
        vendor=VENDOR,
        namespace=namespace,
        custom_image={
            "health_route": "/health",
            "env": {
                "MAX_BATCH_TOKENS": str(MAX_WORKERS * 2048),
                "MAX_CONCURRENT_REQUESTS": "512",
                "MODEL_ID": "/repository",
            },
            "url": "ghcr.io/huggingface/text-embeddings-inference:0.5.0",
        },
        type="protected",
    )
except:
    endpoint = [ie for ie in list_inference_endpoints(namespace=namespace) if ie.name == ENDPOINT_NAME][0]
    print("Loaded endpoint")

>>> %%time
>>> endpoint.wait()

response = endpoint.client.post(
    json={
        "inputs": "This sound track was beautiful! It paints the senery in your mind so well I would recomend it even to people who hate vid. game music!",
        "truncate": True,
    },
    task="feature-extraction",
)
response = np.array(json.loads(response.decode()))
response[0][:20]

>>> embedding_input = "This input will get multiplied" * 10000
>>> print(f"The length of the embedding_input is: {len(embedding_input)}")
>>> response = endpoint.client.post(json={"inputs": embedding_input, "truncate": True}, task="feature-extraction")
>>> response = np.array(json.loads(response.decode()))
>>> response[0][:20]

async def request(document, semaphore):
    # Semaphore guard
    async with semaphore:
        result = await endpoint.async_client.post(
            json={"inputs": document["content"], "truncate": True}, task="feature-extraction"
        )
        result = np.array(json.loads(result.decode()))
        document["embedding"] = result[0]  # Assuming the API's output can be directly assigned
        return document


async def main(documents):
    # Semaphore to limit concurrent requests. Adjust the number as needed.
    semaphore = asyncio.BoundedSemaphore(MAX_WORKERS)

    # Creating a list of tasks
    tasks = [request(document, semaphore) for document in documents]

    # Using tqdm to show progress. It's been integrated into the async loop.
    for f in tqdm(asyncio.as_completed(tasks), total=len(documents)):
        await f

>>> start = time.perf_counter()

>>> # Get embeddings
>>> await main(documents)

>>> # Make sure we got it all
>>> count = 0
>>> for document in documents:
...     if "embedding" in document.keys() and len(document["embedding"]) == 768:
...         count += 1
>>> print(f"Embeddings = {count} documents = {len(documents)}")


>>> # Print elapsed time
>>> elapsed_time = time.perf_counter() - start
>>> minutes, seconds = divmod(elapsed_time, 60)
>>> print(f"{int(minutes)} min {seconds:.2f} sec")

>>> endpoint = endpoint.pause()

>>> print(f"Endpoint Status: {endpoint.status}")

df = pd.DataFrame(documents)
dd = DatasetDict({"train": Dataset.from_pandas(df)})

dd.push_to_hub(repo_id=DATASET_OUT)

>>> print(f'Dataset is at https://huggingface.co./datasets/{who["name"]}/{DATASET_OUT}')

>>> dashboard_url = f"https://ui.endpoints.huggingface.co/{namespace}/endpoints/{ENDPOINT_NAME}"
>>> print(dashboard_url)

>>> input("Hit enter to continue with the notebook")

>>> endpoint = endpoint.delete()

>>> if not endpoint:
...     print("Endpoint deleted successfully")
>>> else:
...     print("Delete Endpoint in manually")