jamba1.5-retriever/JambaRetriever.ipynb (275 lines of code) (raw):

{ "cells": [ { "cell_type": "code", "execution_count": null, "id": "1c907c39-40a2-48a5-8f9f-342c6c300d00", "metadata": {}, "outputs": [], "source": [ "# Import required modules, initialize SageMaker Session and define key parameters for Training and Deployment\n", "# Replace huggingface_token with your token below\n", "\n", "import sagemaker\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "\n", "from sagemaker.huggingface import HuggingFace\n", "from datasets import load_dataset\n", "from sagemaker.huggingface.model import HuggingFaceModel\n", "from sklearn.metrics.pairwise import cosine_similarity\n", "\n", "role = sagemaker.get_execution_role()\n", "sess = sagemaker.Session()\n", "\n", "# Get the default bucket created by SageMaker\n", "default_s3_bucket = sess.default_bucket()\n", "\n", "print(f\"Default S3 bucket: {default_s3_bucket}\")\n", "\n", "model_name = \"ai21labs/Jamba-tiny-dev\"\n", "cache_dir_ds = \"/opt/ml/dataset_cache\"\n", "cache_dir_model = \"/opt/ml/model_cache\"\n", "output_dir = \"/opt/ml/model\"\n", "log_dir = \"/opt/ml/output\"\n", "dataset_name = \"stsb_multi_mt\"\n", "\n", "# Jamba1.5 models are gated models in Huggingface\n", "# You need to generate a HuggingFace User Access Token\n", "# and enable model access via your HuggingFace Account\n", "# https://huggingface.co/docs/hub/en/models-gated\n", "\n", "huggingface_token = \"<Replace_with_your_token>\"\n", "assert huggingface_token != \"<Replace_with_your_token>\", \"Replace with your HuggingFace Token to gain access to gated Jamba1.5 models\"" ] }, { "cell_type": "code", "execution_count": null, "id": "9bc9baf6-90a3-4aa1-b250-31131c794822", "metadata": { "editable": true, "slideshow": { "slide_type": "" }, "tags": [] }, "outputs": [], "source": [ "# Define Hugging Face training estimator\n", "huggingface_estimator = HuggingFace(\n", " entry_point='train.py', # train script provided under /scripts folder\n", " source_dir='./scripts', # Path to the script directory\n", " instance_type='ml.p3.8xlarge',\n", " instance_count=1,\n", " role=role,\n", " transformers_version='4.36.0', # This will get updated via requirements.txt. Transformer support for JAMBA came after 4.39. DLC Training still up to 4.36\n", " pytorch_version='2.1.0',\n", " py_version='py310',\n", " hyperparameters={\n", " 'epochs': 1,\n", " 'train_batch_size': 64,\n", " 'eval_batch_size': 128,\n", " 'learning_rate': 2e-5,\n", " 'model_name': model_name,\n", " 'output_dir': output_dir,\n", " 'log_dir': log_dir,\n", " 'cache_dir_ds': cache_dir_ds,\n", " 'cache_dir_model': cache_dir_model,\n", " 'huggingface_token': huggingface_token,\n", " 'dataset_name': dataset_name\n", " },\n", " enable_sagemaker_metrics=True,\n", " container_log_level=10,\n", " volume_size=1000,\n", " dependencies=['./scripts/requirements.txt'] # needed to install requirements for JAMBA not included in DLC\n", ")\n", "\n", "# Start the fine-tuning job\n", "huggingface_estimator.fit()\n", "\n", "# Save fine-tuned model data location in S3\n", "model_data = huggingface_estimator.model_data" ] }, { "cell_type": "code", "execution_count": 4, "id": "1d9a8431-0c4a-4f02-bcfb-3684c514da95", "metadata": {}, "outputs": [], "source": [ "# Define HuggingFace Model based on the fine tuned model data\n", "huggingface_model = HuggingFaceModel(\n", " model_data=model_data,\n", " source_dir='./scripts',\n", " entry_point='inference.py', # Inference script provided under /scripts folder\n", " role=role,\n", " transformers_version=\"4.37.0\",\n", " pytorch_version=\"2.1.0\",\n", " py_version='py310',\n", " dependencies=['./scripts/requirements.txt'],\n", " env={ # Define environment variables\n", " 'HF_TASK': 'feature-extraction', # Set task to feature-extraction\n", " 'SAGEMAKER_PRELOAD_MODELS': 'true', # Preload model on container startup\n", " 'SAGEMAKER_MODEL_SERVER_TIMEOUT': '180', # Set a high timeout if needed\n", " 'SAGEMAKER_CONTAINER_LOG_LEVEL': 10 # Set log level to DEBUG (most verbose)\n", " }\n", ")" ] }, { "cell_type": "code", "execution_count": null, "id": "ad08cf94-d9da-4d40-b962-74cbb7d9d36e", "metadata": { "editable": true, "slideshow": { "slide_type": "" }, "tags": [] }, "outputs": [], "source": [ "# Define Predictor and start Deployment\n", "predictor = huggingface_model.deploy(\n", " initial_instance_count=1, \n", " instance_type='ml.p3.2xlarge'\n", ")" ] }, { "cell_type": "code", "execution_count": 6, "id": "38ce2d74-9323-4bac-ac52-c60a878f029a", "metadata": { "scrolled": true }, "outputs": [], "source": [ "# Create Sample Test and Receive Sentence Embeddings from fine tuned Model\n", "\n", "Sentence0 = {\n", " \"inputs\": \"ja esta muito tarde\"\n", "}\n", "\n", "Sentence1 = {\n", " \"inputs\": \"I love you\"\n", "}\n", "\n", "Sentence2 = {\n", " \"inputs\": \"can you do my homework tomorrow\"\n", "}\n", "\n", "Sentence3 = {\n", " \"inputs\": \"I'm in love with you\"\n", "}\n", "\n", "SentenceEmbeddings = []\n", "SentenceEmbeddings.append(predictor.predict(data=Sentence0))\n", "SentenceEmbeddings.append(predictor.predict(data=Sentence1))\n", "SentenceEmbeddings.append(predictor.predict(data=Sentence2))\n", "SentenceEmbeddings.append(predictor.predict(data=Sentence3))" ] }, { "cell_type": "code", "execution_count": null, "id": "11a74754-4c59-4cb7-9ae8-5eca6812fdf0", "metadata": {}, "outputs": [], "source": [ "# Function to calculate the distance between embeddings\n", "def compute_cosine_similarity(embeddings1, embeddings2):\n", " # Calculate cosine similarity between two sets of embeddings\n", " return cosine_similarity(embeddings1, embeddings2)\n", "\n", "\n", "# Prepare a function to plot the 2x2 matrix with color coding\n", "def plot_similarity_matrix(similarities, labels, threshold=0.5):\n", " fig, ax = plt.subplots()\n", "\n", " # Create a 2x2 matrix for the results\n", " matrix = np.zeros((2, 2))\n", "\n", " # Fill the matrix based on similarity values and labels\n", " for i in range(2):\n", " for j in range(2):\n", " is_similar = labels[i * 2 + j] # Expected similarity (1 if similar, 0 if dissimilar)\n", " is_good = (similarities[i * 2 + j] >= threshold and is_similar == 1) or (similarities[i * 2 + j] < threshold and is_similar == 0)\n", " matrix[i, j] = similarities[i * 2 + j]\n", "\n", " # Set color: Green for good results, Red for bad results\n", " color = 'green' if is_good else 'red'\n", " ax.text(j, i, f'{matrix[i, j]:.2f}', ha='center', va='center', color='white', fontsize=12, bbox=dict(facecolor=color, alpha=0.7))\n", "\n", " ax.imshow(matrix, cmap='RdYlGn', vmin=0, vmax=1)\n", " ax.set_xticks(np.arange(2))\n", " ax.set_yticks(np.arange(2))\n", " ax.set_xticklabels(['Sentence 1', 'Sentence 3'])\n", " ax.set_yticklabels(['Sentence 0', 'Sentence 2'])\n", "\n", " plt.title('Cosine Similarity Matrix (Green=Good, Red=Bad)')\n", " plt.show()\n", "\n", "\n", "# Expected similarity labels for the test pairs (1 for similar, 0 for dissimilar)\n", "expected_labels = [1, 0, 0, 0] # 1st pair is similar, 2nd pair is dissimilar \n", "\n", "# Initialize similarities\n", "similarities = []\n", "\n", "# Calculate cosine similarity between embeddings\n", "similarities.append(cosine_similarity(SentenceEmbeddings[0].get('embeddings'), SentenceEmbeddings[1].get('embeddings')))\n", "similarities.append(cosine_similarity(SentenceEmbeddings[2].get('embeddings'), SentenceEmbeddings[3].get('embeddings')))\n", "similarities.append(cosine_similarity(SentenceEmbeddings[0].get('embeddings'), SentenceEmbeddings[2].get('embeddings')))\n", "similarities.append(cosine_similarity(SentenceEmbeddings[1].get('embeddings'), SentenceEmbeddings[3].get('embeddings')))\n", "\n", "# Plot the similarity matrix\n", "plot_similarity_matrix(similarities, expected_labels)\n", "print('Similarity Scores: ', similarities)" ] }, { "cell_type": "code", "execution_count": null, "id": "efae36e0-7bc7-4a71-8577-8e1aeda2d0dc", "metadata": {}, "outputs": [], "source": [ "# Delete the model and endpoint when done\n", "predictor.delete_model()\n", "predictor.delete_endpoint()" ] }, { "cell_type": "code", "execution_count": null, "id": "c8b42a49-2e25-4c5c-bb16-7d217f0ab935", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "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.11.9" } }, "nbformat": 4, "nbformat_minor": 5 }