update to allow enforcing pre-query-template

README.md

@@ -87,6 +87,7 @@ Optionally, you can use different models and APIs.
 - `BASE_URL`: The base URL for any OpenAI compatible API, e.g. `https://api-inference.huggingface.co/v1/`, `https://api.openai.com/v1/`.
 - `MODEL`: The model to use for generating the dataset, e.g. `meta-llama/Meta-Llama-3.1-8B-Instruct`, `gpt-4o`.
 - `API_KEY`: The API key to use for the generation API, e.g. `hf_...`, `sk-...`. If not provided, it will default to the provided `HF_TOKEN` environment variable.
+- `MAGPIE_PRE_QUERY_TEMPLATE`: Enforce setting the pre-query template for Magpie generation to either `llama3`, `qwen2`. Not that this is only used if the model is a Qwen or Llama model. If you want to use other model families for chat data generation, feel free to [implement your own pre-query template](https://github.com/argilla-io/distilabel/pull/778/files).
 
 Optionally, you can also push your datasets to Argilla for further curation by setting the following environment variables:
 

src/synthetic_dataset_generator/constants.py

@@ -33,10 +33,20 @@ if BASE_URL != "https://api-inference.huggingface.co/v1/" and len(API_KEYS) == 0
     raise ValueError(
         "API_KEY is not set. Ensure you have set the API_KEY environment variable that has access to the Hugging Face Inference Endpoints."
     )
-
 llama_options = ["llama3", "llama-3", "llama 3"]
 qwen_options = ["qwen2", "qwen-2", "qwen 2"]
-if MODEL.lower() in llama_options:
+if os.getenv("MAGPIE_PRE_QUERY_TEMPLATE"):
+    SFT_AVAILABLE = True
+    passed_pre_query_template = os.getenv("MAGPIE_PRE_QUERY_TEMPLATE")
+    if passed_pre_query_template.lower() in llama_options:
+        MAGPIE_PRE_QUERY_TEMPLATE = "llama3"
+    elif passed_pre_query_template.lower() in qwen_options:
+        MAGPIE_PRE_QUERY_TEMPLATE = "qwen2"
+    else:
+        raise ValueError(
+            f"MAGPIE_PRE_QUERY_TEMPLATE must be either {llama_options} or {qwen_options}."
+        )
+elif MODEL.lower() in llama_options:
     SFT_AVAILABLE = True
     MAGPIE_PRE_QUERY_TEMPLATE = "llama3"
 elif MODEL.lower() in qwen_options:
@@ -49,6 +59,7 @@ else:
     )
     MAGPIE_PRE_QUERY_TEMPLATE = None
 
+
 # Embeddings
 STATIC_EMBEDDING_MODEL = "minishlab/potion-base-8M"