Update README.md

README.md

@@ -2510,20 +2510,29 @@ _coming soon_
 
 ```python
 
-from optimum.neuron import pipeline
+from optimum.neuron import NeuronModelForFeatureExtraction
+from transformers import AutoTokenizer
+import torch
+import torch_neuronx
 
-# Load pipeline from Hugging Face repository
-pipe = pipeline("text-generation", "aws-neuron/bge-base-en-v1-5-seqlen-384-bs-1")
+# Load Model from Hugging Face repository
+model = NeuronModelForFeatureExtraction.from_pretrained("aws-neuron/bge-base-en-v1-5-seqlen-384-bs-1")
+tokenizer = AutoTokenizer.from_pretrained("aws-neuron/bge-base-en-v1-5-seqlen-384-bs-1")
 
-# We use the tokenizer's chat template to format each message - see https://huggingface.co/docs/transformers/main/en/chat_templating
-messages = [
-    {"role": "user", "content": "What is 2+2?"},
-]
-prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-# Run generation
-outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
-print(outputs[0]["generated_text"])
+# sentence input
+inputs = "Hello, my dog is cute"
 
+# Tokenize sentences
+encoded_input = tokenizer(inputs,return_tensors="pt",truncation=True,max_length=model.config.neuron["static_sequence_length"])
+
+# Compute embeddings
+with torch.no_grad():
+    model_output = model(*tuple(encoded_input.values()))
+
+# Perform pooling. In this case, cls pooling.
+sentence_embeddings = model_output[0][:, 0]
+# normalize embeddings
+sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)   
 ```
 
 **input_shapes**