Skip to content

Commit

Permalink
Update task to use system prompt
Browse files Browse the repository at this point in the history
  • Loading branch information
gabrielmbmb committed Jun 3, 2024
1 parent b7cd785 commit 53019bc
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 16 deletions.
30 changes: 22 additions & 8 deletions src/distilabel/steps/tasks/sentence_transformers.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,11 +31,10 @@
GenerationAction = Literal["paraphrase", "semantically-similar", "query"]

POSITIVE_NEGATIVE_PAIR_REGEX = re.compile(
r"## Positive\s+(.*?)\s+(?:## Negative\s+(.*?)\s*)?$",
r"## Positive\s+(.*?)(?:\s+## Negative\s+(.*?))?\s*$",
re.DOTALL,
)


GENERATION_ACTION_SENTENCES: Final[Dict[GenerationAction, str]] = {
"paraphrase": "paraphrase",
"semantically-similar": "be semantically similar to",
Expand Down Expand Up @@ -89,14 +88,27 @@ def inputs(self) -> List[str]:
return ["anchor"]

def format_input(self, input: Dict[str, Any]) -> "ChatType":
action_sentence = GENERATION_ACTION_SENTENCES[self.action]
return [
{
"role": "user",
"content": self._template.render(
anchor=input["anchor"],
action=GENERATION_ACTION_SENTENCES[self.action],
"role": "system",
"content": (
"Your task is to generate a positive and a negative sentence given"
f" an anchor sentence. The positive sentence has to {action_sentence}"
" the anchor sentence, while the negative sentence has to do the opposite."
" You must output only two new sections: `## Positive` and `## Negative`."
)
if self.triplet
else (
"Your task is to generate a positive sentence given an anchor sentence."
f" The positive sentence has to {action_sentence} the anchor sentence."
" You must output only one new section: `## Positive`."
),
}
},
{
"role": "user",
"content": self._template.render(anchor=input["anchor"]),
},
]

@property
Expand All @@ -120,7 +132,9 @@ def format_output(
if self.triplet:
return {
"positive": groups[0].strip(),
"negative": groups[1].strip() if len(groups) > 1 else None,
"negative": groups[1].strip()
if len(groups) > 1 and groups[1] is not None
else None,
}

return {"positive": groups[0].strip()}
Original file line number Diff line number Diff line change
@@ -1,12 +1,4 @@
# Task Description
{% if triplet %}
Your task is to generate a positive and a negative sentence given an anchor sentence. The positive sentence has to {{ action }} the anchor sentence, while the negative sentence is the contrary. You must output only two new sections: `## Positive` and `## Negative`.
{% else %}
Your task is to generate a positive sentence given an anchor sentence. The positive sentence has to {{ action }} the anchor sentence. You must output only one new section: `## Positive`.
{% endif %}

## Anchor

{{ anchor }}

## Positive

0 comments on commit 53019bc

Please sign in to comment.