Update README.md

README.md

@@ -29,7 +29,42 @@ base_model:
 ## How to Get Started 
 
 ```python
-
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
+import torch
+from utils import tokenize_desk, SYSTEM_PROMPT
+
+# Load the mode
+
+
+model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16).to(device)
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+# Define your workspace
+objects = [
+    {"red-cube": [51, 43, 17]},
+    {"black-cube": [44, 58, 17]},
+    {"purple-cube": [74, 59, 17]},
+    {"green-cube": [65, 82, 17]},
+]
+
+# Give a natural language instruction
+instruction = "Throw the red cube on top of the blue cylinder"
+desk, object_height = tokenize_desk(objects)
+final_instruction = SYSTEM_PROMPT.format(object_height=object_height,instruction=instruction,TABLE_MAP=desk)
+chat = [
+    {"role": "user", "content": final_instruction.strip()}
+]
+tokenized_chat = tokenizer.apply_chat_template(chat, tokenize=True, add_generation_prompt=True, use_system_prompt=False, return_tensors="pt")
+# print(len(tokenized_chat[0]))
+generated_ids = model.generate(
+    tokenized_chat.to("cuda"),
+    max_new_tokens=2048,
+    do_sample=False,
+    temperature=0.6,
+)
+# Get the solution
+result = tokenizer.decode(generated_ids[0][tokenized_chat.shape[1]:], skip_special_tokens=True)
+print(result)
 ```
 ### Hardware