File size: 6,378 Bytes
3c9d064 |
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 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 |
import json
from textwrap import dedent
from typing import Any, Dict, List
from distilabel.llms.huggingface import InferenceEndpointsLLM
from distilabel.pipeline import Pipeline
from distilabel.steps import TextGenerationToArgilla
from distilabel.steps.expand import ExpandColumns
from distilabel.steps.generators.data import LoadDataFromDicts
from distilabel.steps.tasks.self_instruct import SelfInstruct
from distilabel.steps.tasks.text_generation import TextGeneration
from distilabel.steps.tasks.typing import ChatType
################################################################################
# Functions to create task prompts
################################################################################
def create_application_instruction(domain: str, examples: List[Dict[str, str]]):
"""Create the instruction for Self-Instruct task."""
system_prompt = dedent(
f"""You are an AI assistant than generates queries around the domain of {domain}.
Your should not expect basic but profound questions from your users.
The queries should reflect a diversxamity of vision and economic positions and political positions.
The queries may know about different methods of {domain}.
The queries can be positioned politically, economically, socially, or practically.
Also take into account the impact of diverse causes on diverse domains."""
)
for example in examples:
question = example["question"]
answer = example["answer"]
system_prompt += f"""\n- Question: {question}\n- Answer: {answer}\n"""
def create_seed_terms(topics: List[str], perspectives: List[str]) -> List[str]:
"""Create seed terms for self intruct to start from."""
return [
f"{topic} from a {perspective} perspective"
for topic in topics
for perspective in perspectives
]
################################################################################
# Define out custom step for the domain expert
################################################################################
class DomainExpert(TextGeneration):
"""A customized task to generate text as a domain expert in the domain of farming and agriculture."""
system_prompt: str
template: str = """This is the the instruction: {instruction}"""
def format_input(self, input: Dict[str, Any]) -> "ChatType":
return [
{
"role": "system",
"content": self.system_prompt,
},
{
"role": "user",
"content": self.template.format(**input),
},
]
################################################################################
# Main script to run the pipeline
################################################################################
if __name__ == "__main__":
import argparse
import json
parser = argparse.ArgumentParser(
description="Run the pipeline to generate domain-specific datasets."
)
parser.add_argument("--hub-token", type=str, help="The Hugging Face API token.")
parser.add_argument("--argilla-api-key", type=str, help="The Argilla API key.")
parser.add_argument("--argilla-api-url", type=str, help="The Argilla API URL.")
parser.add_argument(
"--argilla-dataset-name", type=str, help="The name of the dataset in Argilla."
)
parser.add_argument(
"--seed_data_path",
type=str,
help="The path to the seed data.",
default="seed_data.json",
)
parser.add_argument(
"--endpoint-base-url", type=str, help="The base URL of the inference endpoint."
)
args = parser.parse_args()
# collect our seed data
with open(args.seed_data_path, "r") as f:
seed_data = json.load(f)
topics = seed_data.get("topics", [])
perspectives = seed_data.get("perspectives", [])
domain_expert_prompt = seed_data.get("domain_expert_prompt", "")
examples = seed_data.get("examples", [])
domain_name = seed_data.get("domain_name", "domain")
# Define the task prompts
terms = create_seed_terms(topics=topics, perspectives=perspectives)
application_instruction = create_application_instruction(
domain=domain_name, examples=examples
)
# Define the distilabel pipeline
with Pipeline(domain_name) as pipeline:
load_data = LoadDataFromDicts(
name="load_data",
data=[{"input": term} for term in terms],
batch_size=64,
)
self_instruct = SelfInstruct(
name="self_instruct",
num_instructions=5,
input_batch_size=8,
llm=InferenceEndpointsLLM(
base_url=args.endpoint_base_url,
api_key=args.hub_token,
),
)
expand_instructions = ExpandColumns(
name="expand_columns", columns={"instructions": "instruction"}
)
domain_expert = DomainExpert(
name="domain_expert",
llm=InferenceEndpointsLLM(
base_url=args.endpoint_base_url,
api_key=args.hub_token,
),
input_batch_size=8,
system_prompt=domain_expert_prompt,
)
to_argilla = TextGenerationToArgilla(
name="text_generation_to_argilla",
dataset_name=args.argilla_dataset_name,
dataset_workspace="admin",
api_url=args.argilla_api_url,
api_key=args.argilla_api_key,
)
# Connect up the pipeline
load_data.connect(self_instruct)
self_instruct.connect(expand_instructions)
expand_instructions.connect(domain_expert)
domain_expert.connect(to_argilla)
# Run the pipeline
pipeline.run(
parameters={
"self_instruct": {
"llm": {"api_key": args.hub_token, "base_url": args.endpoint_base_url}
},
"domain_expert": {
"llm": {"api_key": args.hub_token, "base_url": args.endpoint_base_url}
},
"text_generation_to_argilla": {
"dataset_name": args.argilla_dataset_name,
"api_key": args.argilla_api_key,
"api_url": args.argilla_api_url,
},
},
use_cache=False,
)
|