Update README.md

README.md

@@ -157,6 +157,42 @@ predictor.predict(
 
 ```
 
+### Invoke the model on a SageMaker Endpoint ###
+To use MegaBeam-Mistral-7B-300k on a SageMaker endpoint, please try following this example:
+
+```python
+import boto3
+import json
+
+def call_endpoint(text:str, endpoint_name:str):
+    client = boto3.client("sagemaker-runtime")
+
+    parameters = {
+        "max_new_tokens": 450,
+        "do_sample": True,
+        "temperature": 0.7,
+    }
+
+    payload = {"inputs": text, "parameters": parameters}
+
+    response = client.invoke_endpoint(
+        EndpointName=endpoint_name, Body=json.dumps(payload), ContentType="application/json"
+    )
+
+    output = json.loads(response["Body"].read().decode())
+
+    result = output["generated_text"]
+    return result
+
+# please insert your long prompt/document content here
+prompt = """<s>[INST] What are the main challenges to support long contexts for a Large Language Model? [/INST]"""
+
+#print(prompt)
+endpoint_name = "megaBeam-mistral-7b-300k-2024-05-13-14-23-41-219" # please use a valid endpoint name
+result = call_endpoint(prompt, endpoint_name)
+print(result)
+```
+
 
 ## Limitations ##
 Before using the MegaBeam-Mistral-7B-300k model, it is important to perform your own independent assessment, and take measures to ensure that your use would comply with your own specific quality control practices and standards, and that your use would comply with the local rules, laws, regulations, licenses and terms that apply to you, and your content.