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
}