Open In Colab

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,
    aws_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