Interacting with Embeddings deployed in Amazon SageMaker Endpoint with LlamaIndex

An Amazon SageMaker endpoint is a fully managed resource that enables the deployment of machine learning models, for making predictions on new data.

This notebook demonstrates how to interact with Embedding endpoints using SageMakerEmbedding, unlocking additional llamaIndex features. So, It is assumed that an Embedding is deployed on a SageMaker endpoint.

Setting Up

If you’re opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.

%pip install llama-index-embeddings-sagemaker-endpoint

! pip install llama-index

You have to specify the endpoint name to interact with.

ENDPOINT_NAME = "<-YOUR-ENDPOINT-NAME->"

Credentials should be provided to connect to the endpoint. You can either:

• use an AWS profile by specifying the profile_name parameter, if not specified, the default credential profile will be used.
• Pass credentials as parameters (aws_access_key_id, aws_secret_access_key, aws_session_token, region_name).

for more details check this link.

AWS profile name

from llama_index.embeddings.sagemaker_endpoint import SageMakerEmbedding

AWS_ACCESS_KEY_ID = "<-YOUR-AWS-ACCESS-KEY-ID->"
AWS_SECRET_ACCESS_KEY = "<-YOUR-AWS-SECRET-ACCESS-KEY->"
AWS_SESSION_TOKEN = "<-YOUR-AWS-SESSION-TOKEN->"
REGION_NAME = "<-YOUR-ENDPOINT-REGION-NAME->"

embed_model = SageMakerEmbedding(
    endpoint_name=ENDPOINT_NAME,
    aws_access_key_id=AWS_ACCESS_KEY_ID,
    aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
    aws_session_token=AWS_SESSION_TOKEN,
    region_name=REGION_NAME,
)

With credentials:

from llama_index.embeddings.sagemaker_endpoint import SageMakerEmbedding

ENDPOINT_NAME = "<-YOUR-ENDPOINT-NAME->"
PROFILE_NAME = "<-YOUR-PROFILE-NAME->"
embed_model = SageMakerEmbedding(
    endpoint_name=ENDPOINT_NAME, profile_name=PROFILE_NAME
)  # Omit the profile name to use the default profile

Basic Usage

Call get_text_embedding

embeddings = embed_model.get_text_embedding(
    "An Amazon SageMaker endpoint is a fully managed resource that enables the deployment of machine learning models, specifically LLM (Large Language Models), for making predictions on new data."
)

embeddings

[0.021565623581409454,
 0.019147753715515137,
 -0.056700449436903,
 -0.022435730323195457,
 -0.03930153697729111,
 0.0028269595932215452,
 0.13571178913116455,
 -0.01545825693756342,
 0.008395835757255554,
 -0.028859995305538177,
 -0.002176997484639287,
 -0.01331179216504097,
 0.04315004497766495,
 -0.06485316157341003,
 0.05034399777650833,
 0.03711436316370964,
 0.08230344206094742,
 -0.025957593694329262,
 -0.022912224754691124,
 -0.0353832021355629,
 0.057388801127672195,
 0.06676290929317474,
 -0.11050011962652206,
 -0.06368492543697357,
 0.06308538466691971,
 -0.018852291628718376,
 0.03129011392593384,
 -0.0760543942451477,
 -0.03160060569643974,
 0.018145989626646042,
 0.023656774312257767,
 -0.03639936074614525,
 0.017052484676241875,
 -0.02450205385684967,
 0.027692550793290138,
 0.0025261295959353447,
 0.027999257668852806,
 -0.01860380358994007,
 0.037128496915102005,
 -0.08015701919794083,
 0.06850703060626984,
 -0.008332699537277222,
 0.06535708159208298,
 0.0328301340341568,
 0.006906852591782808,
 0.02144588902592659,
 0.016819249838590622,
 0.065149687230587,
 -0.044132206588983536,
 -0.06759986281394958,
 -0.06730470061302185,
 0.04957125708460808,
 0.06852839142084122,
 0.025557655841112137,
 0.02020672708749771,
 0.010702744126319885,
 -0.08458871394395828,
 0.015973811969161034,
 -0.04692167788743973,
 0.006929733324795961,
 -0.07795898616313934,
 -0.0068844580091536045,
 0.00868065282702446,
 0.05663231015205383,
 0.0727660208940506,
 0.025757266208529472,
 -0.005326052196323872,
 -0.01100397203117609,
 -0.12615707516670227,
 -0.05008727312088013,
 0.011879516765475273,
 0.05348730832338333,
 0.012098579667508602,
 -0.09577080607414246,
 0.10378099977970123,
 0.032600581645965576,
 0.030491407960653305,
 -0.08457385003566742,
 0.008462479338049889,
 0.0191865935921669,
 -0.07925017923116684,
 -0.014723804779350758,
 -0.030580421909689903,
 -0.042236607521772385,
 -0.05877465009689331,
 0.010675436817109585,
 0.032128408551216125,
 -0.0562652088701725,
 0.02557305432856083,
 -0.04288778454065323,
 0.06932021677494049,
 -0.0056538935750722885,
 -0.034259598702192307,
 -0.03814560920000076,
 -0.050656333565711975,
 -0.07511229068040848,
 -0.0727902203798294,
 0.057129207998514175,
 0.0038894452154636383,
 -0.026942580938339233,
 0.0033212697599083185,
 -0.0017185898032039404,
 0.09014856070280075,
 -0.050106242299079895,
 -0.03977276012301445,
 0.09139455854892731,
 -0.047016214579343796,
 -0.08531513810157776,
 0.028743667528033257,
 -0.09533917158842087,
 -0.0395311564207077,
 0.05482438951730728,
 0.08327925205230713,
 0.0976642593741417,
 0.009845257736742496,
 -0.04681103676557541,
 -0.005461001303046942,
 -0.008060128428041935,
 0.03265204653143883,
 -0.030966755002737045,
 0.08599787205457687,
 -0.10304504632949829,
 -0.026105130091309547,
 -0.059405822306871414,
 0.0397467203438282,
 -0.10144900530576706,
 -0.08908584713935852,
 -0.002865528455004096,
 0.018968375399708748,
 -0.015612486749887466,
 -0.02899557165801525,
 0.014865572564303875,
 0.02968095801770687,
 0.02684851363301277,
 0.05083318054676056,
 0.02924215793609619,
 0.01617436856031418,
 -0.12093227356672287,
 -0.03196290135383606,
 -0.06529344618320465,
 -0.026738135144114494,
 -0.006192725617438555,
 -0.0158796738833189,
 0.03231367841362953,
 0.0864284336566925,
 0.009570070542395115,
 0.0716225802898407,
 -0.07572516798973083,
 0.05317925661802292,
 0.06873224675655365,
 -0.0053744143806397915,
 0.07709360867738724,
 0.04732672870159149,
 0.030548326671123505,
 0.0027018985711038113,
 0.008613159880042076,
 -0.027862899005413055,
 -0.047938231378793716,
 -0.0024136374704539776,
 0.10075152665376663,
 -0.05304369330406189,
 -0.002246642718091607,
 -0.029105287045240402,
 -0.06296665966510773,
 -0.014757883735001087,
 -0.052865151315927505,
 0.006513366475701332,
 0.06221649795770645,
 -0.011616945266723633,
 -0.009876491501927376,
 -0.09408566355705261,
 0.031154463067650795,
 -0.08682437986135483,
 -0.06839584559202194,
 -0.004306801129132509,
 -0.04914987087249756,
 -0.06964779645204544,
 -0.07772491127252579,
 0.020129483193159103,
 0.046311721205711365,
 -0.05257619544863701,
 -0.019545624032616615,
 -0.08561691641807556,
 -0.0333973690867424,
 -0.07877060025930405,
 0.02237010933458805,
 0.04267224296927452,
 -0.02596845105290413,
 -0.07537004351615906,
 -0.09591358155012131,
 -0.02983330562710762,
 -0.013941322453320026,
 -0.011087455786764622,
 -0.012412399053573608,
 -0.02150963619351387,
 0.07786650210618973,
 0.04773382097482681,
 0.019922567531466484,
 0.07738719880580902,
 0.05013704299926758,
 -0.09917087852954865,
 0.06646912544965744,
 -0.029166480526328087,
 -0.01662266068160534,
 0.052809495478868484,
 0.017951449379324913,
 -0.08008625358343124,
 0.07121209800243378,
 -0.02369554713368416,
 -0.05586501955986023,
 0.0498092845082283,
 -0.01365865021944046,
 -0.02708507515490055,
 0.08256032317876816,
 0.025594351813197136,
 -0.061894018203020096,
 0.012498111464083195,
 0.11298499256372452,
 -0.002179534872993827,
 0.04463203623890877,
 0.010685713030397892,
 0.03099050745368004,
 0.05935501679778099,
 0.058455295860767365,
 -0.02524683251976967,
 0.07769811153411865,
 0.03721567243337631,
 -0.05615007132291794,
 0.01681063137948513,
 0.04380310699343681,
 -0.01826358400285244,
 -0.008245888166129589,
 -0.05828167498111725,
 -0.0558970607817173,
 0.0023540062829852104,
 -0.02012726105749607,
 -0.0008245687349699438,
 0.00417731748893857,
 0.10165227949619293,
 -0.04578690975904465,
 -0.026630911976099014,
 -0.0011021343525499105,
 0.022034477442502975,
 0.028991390019655228,
 0.05962913855910301,
 -0.0630992129445076,
 -0.06605151295661926,
 -0.04041984677314758,
 0.024546820670366287,
 0.09267142415046692,
 0.08537988364696503,
 -0.004139745142310858,
 0.005737761966884136,
 0.009916293434798717,
 0.06632582098245621,
 0.06800509989261627,
 0.055296897888183594,
 0.02332977205514908,
 -0.0944223552942276,
 0.017701460048556328,
 -0.08671680837869644,
 -0.03739886358380318,
 -0.05368325486779213,
 0.024832075461745262,
 -0.0017475909553468227,
 -0.01282496377825737,
 -0.02055644989013672,
 -0.03824482858181,
 -0.0029880295041948557,
 -0.06276058405637741,
 -0.0031180321238934994,
 0.018192732706665993,
 -0.03707187995314598,
 0.021848147734999657,
 0.02335541695356369,
 -0.01648322306573391,
 0.1435147374868393,
 0.09290475398302078,
 0.006393723655492067,
 -0.04305407404899597,
 0.07931894063949585,
 0.027527527883648872,
 0.014633690007030964,
 -0.0010841790353879333,
 -0.08786225318908691,
 -0.058670710772275925,
 -0.01426162850111723,
 0.0070671215653419495,
 -0.006778248585760593,
 -0.11174825578927994,
 -0.0030391064938157797,
 -0.04667389392852783,
 0.003977365326136351,
 0.044250018894672394,
 0.07857362180948257,
 0.02811703085899353,
 -0.03371573984622955,
 0.04170072823762894,
 0.04904920980334282,
 0.043538808822631836,
 -0.03201114386320114,
 -0.10568009316921234,
 0.054847411811351776,
 0.04661855846643448,
 -0.014991141855716705,
 0.05715333670377731,
 0.09562608599662781,
 0.07444649189710617,
 0.05350692570209503,
 -0.022917866706848145,
 0.05435822904109955,
 -0.05491911619901657,
 -0.02363002300262451,
 -0.006824106443673372,
 0.03266216069459915,
 0.11281733214855194,
 -0.03085748292505741,
 0.08682505786418915,
 -0.024763137102127075,
 -0.04231179505586624,
 0.0027543380856513977,
 -0.05821159482002258,
 0.011324281804263592,
 0.004590475000441074,
 0.024966349825263023,
 -0.01651289500296116,
 -0.0320640504360199,
 0.02205571159720421,
 0.01905834674835205,
 -0.006711882073432207,
 -0.007124985568225384,
 0.056979432702064514,
 0.013075708411633968,
 -0.013979261741042137,
 -0.11352594941854477,
 -0.03491680696606636,
 0.013942153193056583,
 0.043169185519218445,
 0.0068232291378080845,
 -0.0008051415206864476,
 0.006580108776688576,
 -0.025622187182307243,
 -0.05044765770435333,
 0.05245169252157211,
 -0.0641460195183754,
 -0.055498600006103516,
 -0.007446881849318743,
 -5.276017327560112e-05,
 -0.029843736439943314,
 0.02432435005903244,
 0.06417108327150345,
 -0.06314094364643097,
 0.04182494059205055,
 0.042477674782276154,
 -0.012201766483485699,
 0.09175929427146912,
 -0.021128077059984207,
 -0.03640036657452583,
 0.08847282081842422,
 0.11529214680194855,
 -0.026308033615350723,
 -0.046475768089294434,
 0.04068070650100708,
 0.01918831467628479,
 -0.007530891802161932,
 -0.07138030230998993,
 -0.040803324431180954,
 0.032108619809150696,
 -0.0330362431704998,
 -0.01749657653272152,
 0.034432828426361084,
 0.048382822424173355,
 -0.011264889501035213,
 -0.029528144747018814,
 0.006507426034659147,
 -0.08372807502746582,
 0.10181766003370285,
 -0.041403524577617645,
 -0.0535278283059597,
 -0.03146831691265106,
 0.006305839400738478,
 0.020719777792692184,
 0.03841579705476761,
 0.0018788420129567385]

Call get_text_embedding_batch

embeddings = embed_model.get_text_embedding_batch(
    [
        "An Amazon SageMaker endpoint is a fully managed resource that enables the deployment of machine learning models",
        "Sagemaker is integrated with llamaIndex",
    ]
)

len(embeddings)

2