raffle.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pSVsjsxJYlef"
      },
      "source": [
        "# Raffle ticket writer/picker with DataStax Astra DB vector database and RAGStack\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ih5ZaynCut9V"
      },
      "source": [
        "This notebook takes you through the steps to build a simple application GenAI application to store (and draw) raffle winners. The idea is to be able to enter the names of conference attendees one-at-a-time. When the raffle is over, a random phrase must be entered to generate the winning embedding. The results are returned in the order of their similarity to the random phrase.\n",
        "\n",
        "Requirements:\n",
        " - Access to the [all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2) 384 dimensional sentence transformer on [Hugging Face](https://huggingface.co/).\n",
        " - A free account and vector database with [Astra DB](https://astra.datastax.com/).\n",
        "     - Sufficient resources for this notebook to build a collection named `raffle_data` with a:\n",
        "         - 384 dimensional vector structure.\n",
        "         - cosine metric.\n",
        " - An API endpoint for Astra DB.\n",
        " - An access token for Astra DB.\n",
        "\n",
        "*Note: This notebook will create the `raffle_data` collection if it does not exist.*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q-aQF4kXYW9d"
      },
      "source": [
        "## Install the DataStax RAGStack and SentenceTransformers libraries."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tlC0ANzhXaPx"
      },
      "outputs": [],
      "source": [
        "# @title\n",
        "!pip install astrapy sentence-transformers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WT48XZXLY3mI"
      },
      "source": [
        "## Library Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HU56dcCnY3RN"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "from getpass import getpass\n",
        "from astrapy import DataAPIClient\n",
        "from sentence_transformers import SentenceTransformer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2KeQGuEVaQrq"
      },
      "source": [
        "## Environment Variables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4hHxORm9aTJ2"
      },
      "outputs": [],
      "source": [
        "ASTRA_DB_APPLICATION_TOKEN = getpass('Your Astra DB Token (\"AstraCS:...\"): ')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5inR_IwEaYpH"
      },
      "outputs": [],
      "source": [
        "ASTRA_DB_API_ENDPOINT = input('Your Astra DB API endpoint: ')\n",
        "ASTRA_DB_NAMESPACE='default_keyspace'\n",
        "ASTRA_DB_COLLECTION_NAME='raffle_data'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rYeoodKwaOoX"
      },
      "source": [
        "## Connect to Astra DB"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CM7ccMeSbDIG"
      },
      "outputs": [],
      "source": [
        "client = DataAPIClient(ASTRA_DB_APPLICATION_TOKEN)\n",
        "database = client.get_database_by_api_endpoint(ASTRA_DB_API_ENDPOINT)\n",
        "collection = database.create_collection(ASTRA_DB_COLLECTION_NAME, dimension=384, metric=\"cosine\", check_exists=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ItXH0z1Gdv3p"
      },
      "source": [
        "## Initialize Sentence Transformer \"all-MiniLM-L6-v2\" model locally"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p_dCORradl9L"
      },
      "outputs": [],
      "source": [
        "# initialize the all-MiniLM-L6-v2 model locally\n",
        "model = SentenceTransformer('all-MiniLM-L6-v2')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4iXV4gbIdNWE"
      },
      "source": [
        "## Enter attendee's name and email\n",
        "A vector embedding will be generated for the attendee's name and email address, and it will be stored in Astra DB."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "itLLsFqLeDo5"
      },
      "outputs": [],
      "source": [
        "name = input(\"Enter attendee's name: \")\n",
        "email = input(\"Enter attendee's email address: \")\n",
        "\n",
        "vector_embedding = model.encode(name + email)\n",
        "\n",
        "collection.insert_one({ \"name\": name, \"email\": email }, vector=vector_embedding)\n",
        "\n",
        "# show vector embedding\n",
        "print(vector_embedding)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S5awoSXAgeGh"
      },
      "source": [
        "## Draw winning names"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j6VcUtUhg3ZP"
      },
      "outputs": [],
      "source": [
        "# Generate a vector embedding based on a random string\n",
        "search_string = input(\"Enter a phrase to generate an embedding: \")\n",
        "winner_embedding = model.encode(search_string)\n",
        "print(winner_embedding.tolist())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# run a vector search on this \"random\" embedding, and the top result is your winner!\n",
        "results = collection.vector_find(winner_embedding.tolist(), limit=9, include_similarity=1, fields={\"name\",\"email\"})\n",
        "print(str(results).replace(\"}, {\",\"},\\n{\"))"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.1"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}