Spaces:
Sleeping
Sleeping
File size: 4,595 Bytes
712d204 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
from simulator_env import SwarmAgent
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import textwrap
import re
# Download only the behavior-tree model shard
model_path = hf_hub_download(
repo_id="Inventors-Hub/SwarmChat-models",
repo_type="model",
filename="Falcon3-10B-Instruct-BehaviorTree-3epochs.Q4_K_M.gguf",
)
# llm = Llama(model_path=model_path, n_ctx=1024*4)
llm = Llama(
model_path=model_path,
n_ctx=1024*4, # down from 4096
low_vram=True, # llama.cpp low-vram mode
f16_kv=True, # half-precision kv cache
use_mmap=True, # mmap file
use_mlock=False,
)
def call_behaviors() -> dict:
behavior_dict = {}
for name, attribute in SwarmAgent.__dict__.items():
if callable(attribute) and not name.startswith("_") \
and not name.startswith("update") and not name.startswith("obstacle"):
doc = attribute.__doc__
if doc is not None:
# Dedent, strip, and join into one line by replacing newlines and tabs
cleaned_doc = " ".join(textwrap.dedent(doc).strip().split())
else:
cleaned_doc = ""
behavior_dict[name] = cleaned_doc
return behavior_dict
def extract_behavior_tree(response: str) -> str:
"""
Extracts an XML behavior tree from the given response text.
Looks for a block of XML enclosed in <root...</root> tags.
"""
pattern = re.compile(r'(<root.*?</root>)', re.DOTALL)
match = pattern.search(response)
if match:
return match.group(1).strip()
else:
# If no valid XML block is found, return the original response.
return response.strip()
def save_behavior_tree(tree_xml: str, file_name: str = "tree.xml") -> None:
"""
Saves the behavior tree XML to a file.
"""
with open(file_name, "w", encoding="utf-8") as f:
f.write(tree_xml)
def construct_prompt(prompt: str, prompt_type: str="one") -> str:
behaviors = call_behaviors()
behaviors_text = "\n".join(f"{name}: {doc}" for name, doc in behaviors.items())
plan_prompt = f"""
<s>
<<SYS>>You are a helpful, respectful, and honest AI assistant. Your task is to generate well-structured XML code for behavior trees based on the provided instructions.<</SYS>>
INSTRUCTIONS: It is CRITICAL to use only the following behaviors structured as a dictionary: {behaviors_text} to construct behavior tree in XML format for the following command. Including any behavior that is not in the provided dictionary can result in damage to the agents and potentially humans, therefore you are not allowed to do so. AVOID AT ALL COSTS.
USER COMMAND: generate behavior tree to "{prompt}". Take a step back and think deeply about the behavior you need for this command. Consider the XML structure and the behaviors you use.
The output MUST follow this XML structure exactly, including:
- A root element with <root BTCPP_format and main_tree_to_execute attributes.
- A <BehaviorTree> element with an inner structure of Sequences, Fallback, Conditions, and Actions.
- A <TreeNodesModel> section listing all node models.
- No additional text or commentary outside the XML.
Output only the XML behavior tree without extra text.
OUTPUT:
"""
if prompt_type == "zero":
return plan_prompt
elif prompt_type == "one":
path = "Prompt_One_shot.txt"
with open(path, "r", encoding="utf-8") as file:
file_content = file.read()
return f"{file_content} {plan_prompt}"
elif prompt_type == "two":
path = "Prompt_Two_Shot.txt"
with open(path, "r", encoding="utf-8") as file:
file_content = file.read()
return f"{file_content} {plan_prompt}"
else:
raise ValueError("Unknown prompt type provided.")
def generate_behavior_tree(task_prompt: str) -> str:
prompt = construct_prompt(task_prompt)
print("\n\n",prompt,"\n\n")
output = llm(
prompt,
temperature=0,
max_tokens=1024,
top_p=0.95,
top_k=50,
repeat_penalty=1.1
)
response = output.get("choices", [{}])[0].get("text", "").strip()
tree_xml = extract_behavior_tree(response)
save_behavior_tree(tree_xml)
print("\n response: \n", response)
return tree_xml
# Example usage:
if __name__ == "__main__":
task = "Generate a behavior tree to just form a line."
response = generate_behavior_tree(task)
print("Generated behavior tree response:")
print(response)
|