Skip to content

Commit

Permalink
fix: export dynamic is backend order agnostic
Browse files Browse the repository at this point in the history
  • Loading branch information
@baptiste-olivier
baptiste-olivier committed Nov 14, 2024
1 parent 4261aa6 commit a4abd8f
Show file tree
Hide file tree
Showing 2 changed files with 76 additions and 46 deletions.
98 changes: 64 additions & 34 deletions src/kili/llm/services/export/dynamic.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
CHAT_ITEMS_NEEDED_FIELDS = [
"id",
"content",
"createdAt",
"modelId",
"parentId",
"role",
Expand Down Expand Up @@ -123,41 +122,72 @@ def _init_round(self, context):

def _build_rounds(self, chat_items, annotations, json_interface):
"""A round is composed of a prompt with n pre-prompts and n completions."""
dict_chat_items = {}
for chat_item in chat_items:
if dict_chat_items.get(chat_item["parentId"]) is None:
dict_chat_items[chat_item["parentId"]] = []
dict_chat_items[chat_item["parentId"]].append(chat_item)
rounds = []
parent_target = None
has_children = True
current_round = self._init_round([])
for chat_item in chat_items:
role = chat_item["role"].lower() if chat_item["role"] else None
if role == "user" or role == "system":
if current_round["prompt"] is not None:
rounds.append(current_round)
new_context = (
current_round["context"]
+ current_round["pre_prompts"]
+ [
current_round["prompt"],
self._get_round_winner(
current_round["completion"],
current_round["annotations"],
json_interface,
),
]
)
current_round = self._init_round(new_context)

if role == "user":
current_round["prompt"] = chat_item
elif role == "system":
current_round["pre_prompts"].append(chat_item)
elif role == "assistant":
current_round["completion"].append(chat_item)
else:
raise ValueError(f"Role {chat_item['role']} not supported")
current_round["annotations"] += [
annotation
for annotation in annotations
if annotation["chatItemId"] == chat_item["id"]
]
rounds.append(current_round)

while has_children:
node = dict_chat_items[parent_target][0]
if node["role"].lower() == "system":
current_round["pre_prompts"].append(node)
parent_target = node["id"]
current_round["annotations"] += [
annotation
for annotation in annotations
if annotation["chatItemId"] == node["id"]
]
continue

if node["role"].lower() == "user":
current_round["prompt"] = node
parent_target = node["id"]
current_round["annotations"] += [
annotation
for annotation in annotations
if annotation["chatItemId"] == node["id"]
]
continue

if node["role"].lower() == "assistant":
has_children = False
if dict_chat_items.get(parent_target) is None:
continue
for chat_item in dict_chat_items[parent_target]:
current_round["completion"].append(chat_item)
current_round["annotations"] += [
annotation
for annotation in annotations
if annotation["chatItemId"] == chat_item["id"]
]
if not has_children and dict_chat_items.get(chat_item["id"]) is not None:
has_children = True
parent_target = chat_item["id"]

rounds.append(current_round)
new_context = (
current_round["context"]
+ current_round["pre_prompts"]
+ [
current_round["prompt"],
self._get_round_winner(
current_round["completion"],
current_round["annotations"],
json_interface,
),
]
)
current_round = self._init_round(new_context)
continue

raise ValueError(f"Role {node['role']} not supported")
if current_round["prompt"] is not None:
rounds.append(current_round)
return rounds


Expand Down
24 changes: 12 additions & 12 deletions tests/unit/llm/services/export/test_dynamic.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -111,12 +111,12 @@
},
"chatItems": [
{
"id": "cm2u6kgcc001aj7ja1stsbrvu",
"content": "You are a helpful assistant",
"createdAt": "2024-08-06T12:28:52.170Z",
"modelId": None,
"parentId": None,
"role": "SYSTEM",
"id": "clziefeoe003m7tc976xwbh58",
"content": "Turtles are reptiles known for their protective shells, which act as both home and armor. They are slow-moving on land but can be agile in water, with many species being excellent swimmers. Turtles are omnivorous, feeding on plants, insects, and small animals. They are cold-blooded and rely on external heat sources to regulate their body temperature. Turtles are long-lived, with some species living over 100 years. They are found on every continent except Antarctica, inhabiting a variety of environments including oceans, freshwater lakes, and even deserts. Many turtle species are endangered due to habitat loss and other factors.",
"createdAt": "2024-08-06T12:30:52.430Z",
"modelId": "clzief6pr003c7tc99680e8yj",
"parentId": "clziefeh6003k7tc99abderkk",
"role": "ASSISTANT",
},
{
"id": "clziefeh6003k7tc99abderkk",
Expand All @@ -127,12 +127,12 @@
"role": "USER",
},
{
"id": "clziefeoe003m7tc976xwbh58",
"content": "Turtles are reptiles known for their protective shells, which act as both home and armor. They are slow-moving on land but can be agile in water, with many species being excellent swimmers. Turtles are omnivorous, feeding on plants, insects, and small animals. They are cold-blooded and rely on external heat sources to regulate their body temperature. Turtles are long-lived, with some species living over 100 years. They are found on every continent except Antarctica, inhabiting a variety of environments including oceans, freshwater lakes, and even deserts. Many turtle species are endangered due to habitat loss and other factors.",
"createdAt": "2024-08-06T12:30:52.430Z",
"modelId": "clzief6pr003c7tc99680e8yj",
"parentId": "clziefeh6003k7tc99abderkk",
"role": "ASSISTANT",
"id": "cm2u6kgcc001aj7ja1stsbrvu",
"content": "You are a helpful assistant",
"createdAt": "2024-08-06T12:28:52.170Z",
"modelId": None,
"parentId": None,
"role": "SYSTEM",
},
{
"id": "clziefepk003n7tc9fyx49vei",
Expand Down

0 comments on commit a4abd8f

Please sign in to comment.