|
|
from typing import Dict, List, Optional, Tuple |
|
|
|
|
|
from swift.llm.template import split_str_parts_by |
|
|
|
|
|
|
|
|
def calculate_loss_scale(query: str, |
|
|
response: str, |
|
|
response_loss_scale_map: Dict[str, list], |
|
|
query_loss_scale_map: Optional[Dict[str, list]] = None) -> Tuple[List[str], List[float]]: |
|
|
"""Calculate the loss scale by splitting the agent response. |
|
|
|
|
|
This algorithm comes from paper: https://arxiv.org/pdf/2309.00986.pdf |
|
|
|
|
|
Agent response format: |
|
|
|
|
|
```text |
|
|
Thought: you should always think about what to do |
|
|
Action: the action to take, should be one of the above tools[fire_recognition, |
|
|
fire_alert, call_police, call_fireman] |
|
|
Action Input: the input to the action |
|
|
Observation: the result of the action |
|
|
... (this Thought/Action/Action Input/Observation can be repeated zero or more times) |
|
|
Thought: I now know the final answer |
|
|
Final Answer: the final answer to the original input question |
|
|
``` |
|
|
Returns: |
|
|
A tuple of agent response parts and their weights. |
|
|
""" |
|
|
|
|
|
if query_loss_scale_map is not None: |
|
|
for key in query_loss_scale_map.keys(): |
|
|
if key in query: |
|
|
if isinstance(query_loss_scale_map[key], (float, int)): |
|
|
query_loss_scale_map[key] = [query_loss_scale_map[key]] |
|
|
loss_scale_value = query_loss_scale_map[key][0] |
|
|
return [response], [float(loss_scale_value)] |
|
|
delimiters = [k for k, v in response_loss_scale_map.items() if len(v) == 2] |
|
|
if delimiters: |
|
|
agent_parts = split_str_parts_by(response, delimiters) |
|
|
else: |
|
|
regex_delimiters = [k for k, v in response_loss_scale_map.items() if len(v) == 1] |
|
|
agent_parts = split_str_parts_by(response, regex_delimiters, regex_mode=True) |
|
|
weights = [] |
|
|
agent_content = [] |
|
|
for c in agent_parts: |
|
|
if c['key'] in response_loss_scale_map: |
|
|
loss_scale = response_loss_scale_map[c['key']] |
|
|
assert len(loss_scale) in {1, 2}, f'loss_scale: {loss_scale}' |
|
|
if len(loss_scale) == 1: |
|
|
weights += loss_scale |
|
|
agent_content.append(c['content']) |
|
|
else: |
|
|
weights += loss_scale |
|
|
agent_content += [c['key'], c['content']] |
|
|
else: |
|
|
weights.append(1.) |
|
|
agent_content.append(c['content']) |
|
|
return agent_content, weights |
|
|
|
