russian prompt and augmentation generation

autointent/generation/utterances/basic/chat_template.py

@@ -132,3 +132,119 @@ def __call__(self, intent_data: Intent, n_examples: int) -> list[Message]:
                 f"Please generate {n_examples} more examples for the provided intent class.\n",
             ),
         ]
+
+class SynthesizerChatTemplateRussian(BaseSynthesizer):
+    """Russian language template for generating additional intent examples"""
+
+    __messages: ClassVar[list[Message]] = [
+        Message(
+            role=Role.USER,
+            content=(
+                "Вам будет предоставлен набор примеров высказываний и название общей темы (интент). "
+                "Ваша задача - сгенерировать дополнительные примеры, соответствующие этому интенту.\n\n"
+                "Правила:\n"
+                "- Можно менять значения слотов в похожих высказываниях\n"
+                "- Можно создавать совершенно другие формулировки для того же интента\n"
+                "- Если название интента отсутствует, определите его из примеров\n"
+                "- Если примеры отсутствуют, используйте только название интента\n"
+                "{extra_instructions}\n\n"
+                "Название интента: заказ_пиццы\n\n"
+                "Примеры высказываний:\n"
+                "1. Хочу заказать большую пиццу с пепперони\n"
+                "2. Можно среднюю пиццу с сыром и оливками?\n"
+                "3. Привезите маленькую вегетарианскую пиццу по моему адресу\n\n"
+                "Пожалуйста, сгенерируйте еще 3 примера для этого интента."
+            ),
+        ),
+        Message(
+            role=Role.ASSISTANT,
+            content=(
+                "1. Мне нужна большая пицца Маргарита\n"
+                "2. Можно гавайскую пиццу среднего размера с дополнительным ананасом?\n"
+                "3. Доставьте маленькую пиццу с курицей барбекю на дом"
+            ),
+        ),
+        Message(
+            role=Role.USER,
+            content=(
+                "Название интента: бронирование_отеля\n\n"
+                "Примеры высказываний:\n"
+                "1. Нужно забронировать номер в Москве на две ночи\n\n"
+                "Пожалуйста, сгенерируйте еще 2 примера для этого интента."
+            ),
+        ),
+        Message(
+            role=Role.ASSISTANT,
+            content=(
+                "1. Забронируйте люкс в Санкт-Петербурге на выходные\n"
+                "2. Ищу номер с видом на море в Сочи"
+            ),
+        ),
+        Message(
+            role=Role.USER,
+            content=(
+                "Название интента:\n\n"
+                "Примеры высказываний:\n"
+                "1. Какая сегодня погода?\n\n"
+                "Пожалуйста, сгенерируйте еще 2 примера для этого интента."
+            ),
+        ),
+        Message(
+            role=Role.ASSISTANT,
+            content=(
+                "1. Какой прогноз на завтра?\n"
+                "2. Будет ли дождь в субботу?"
+            ),
+        ),
+        Message(
+            role=Role.USER,
+            content=(
+                "Название интента: запись_на_прием\n\n"
+                "Примеры высказываний:\n\n"
+                "Пожалуйста, сгенерируйте еще 3 примера для этого интента."
+            ),
+        ),
+        Message(
+            role=Role.ASSISTANT,
+            content=(
+                "1. Нужно записаться к врачу на следующую неделю\n"
+                "2. Хочу назначить встречу с парикмахером на пятницу\n"
+                "3. Свободно ли время в пятницу утром для визита?"
+            ),
+        ),
+    ]
+
+    def __init__(
+        self,
+        dataset: Dataset,
+        split: str,
+        extra_instructions: str | None = None,
+        max_sample_utterances: int | None = None,
+    ) -> None:
+        if extra_instructions is None:
+            extra_instructions = ""
+
+        self._messages = deepcopy(self.__messages)
+
+        msg = self._messages[0]
+        msg["content"] = msg["content"].format(extra_instructions=extra_instructions)
+
+        self.dataset = dataset
+        self.split = split
+        self.max_sample_utterances = max_sample_utterances
+
+    def __call__(self, intent_data: Intent, n_examples: int) -> list[Message]:
+        filtered_split = self.dataset[self.split].filter(lambda sample: sample[Dataset.label_feature] == intent_data.id)
+        sample_utterances = filtered_split[Dataset.utterance_feature]
+        if self.max_sample_utterances is not None:
+            sample_utterances = random.sample(sample_utterances, k=self.max_sample_utterances)
+
+        return [
+            *self._messages,
+            Message(
+                role=Role.USER,
+                content=f"Название интента: {intent_data.name}\n\n"
+                f"Примеры высказываний:\n{sample_utterances}\n\n"
+                f"Пожалуйста, сгенерируйте {n_examples} дополнительных примеров для этого интента.\n",
+            ),
+        ]

scripts/data/augment_and_split_dataset.py

@@ -0,0 +1,114 @@
+import json
+import random
+from argparse import ArgumentParser
+from pathlib import Path
+from typing import List
+from collections import defaultdict
+from autointent import Dataset
+from autointent.generation.utterances.basic.chat_template import SynthesizerChatTemplateRussian
+from autointent.generation.utterances.basic.utterance_generator import UtteranceGenerator
+from autointent.generation.utterances.generator import Generator
+
+def process_utterances(generated: List[str]) -> List[str]:
+    processed = []
+    for ut in generated:
+        if "', '" in ut or "',\n" in ut:
+            clean_ut = ut.replace("[", "").replace("]", "").replace("'", "")
+            split_ut = [u.strip() for u in clean_ut.split(", ") if u.strip()]
+            processed.extend(split_ut)
+        else:
+            processed.append(ut.strip())
+    return processed
+
+def main():
+    parser = ArgumentParser()
+    parser.add_argument("--input-path", type=str, required=True, help="Path to few-shot dataset")
+    parser.add_argument("--output-dir", type=str, required=True, help="Directory to save generated datasets")
+    parser.add_argument("--n-augment", type=int, required=True, 
+                      help="Max number of augmented examples per class to generate")
+    parser.add_argument("--split-numbers", type=int, nargs="+", required=True,
+                      help="List of example counts to split into (e.g. 1 2 3 5)")
+    parser.add_argument("--max-attempts", type=int, default=5,
+                      help="Max generation attempts per class")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    args = parser.parse_args()
+
+    random.seed(args.seed)
+    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+
+    dataset = Dataset.from_json(args.input_path)
+    template = SynthesizerChatTemplateRussian(dataset, split="train")
+    generator = UtteranceGenerator(Generator(), template)
+
+    augmented_samples = []
+    for intent in dataset.intents:
+        print(f"\nProcessing intent: {intent.name} (ID: {intent.id})")
+
+        valid_utterances = []
+        attempts = 0
+
+        while len(valid_utterances) < args.n_augment and attempts < args.max_attempts:
+            needed = args.n_augment - len(valid_utterances)
+            generated = generator(intent_data=intent, n_generations=needed)
+
+            processed = process_utterances(generated)
+            current_valid = [
+                ut for ut in processed 
+                if ut and len(ut.split()) > 2
+            ]
+            valid_utterances.extend(current_valid)
+
+            print(f"Attempt {attempts+1}: "
+                  f"Generated {len(current_valid)} valid, "
+                  f"Total {len(valid_utterances)}/{args.n_augment}")
+            attempts += 1
+
+        if len(valid_utterances) < args.n_augment:
+            raise RuntimeError(
+                f"Failed to generate {args.n_augment} examples for "
+                f"{intent.name} after {args.max_attempts} attempts"
+            )
+
+        augmented_samples.extend([
+            {"utterance": ut, "label": intent.id} 
+            for ut in valid_utterances[:args.n_augment]
+        ])
+
+    raw_augmented_path = Path(args.output_dir) / "raw_augmented_samples.json"
+    with open(raw_augmented_path, "w", encoding="utf-8") as f:
+        json.dump({
+            "intents": [{"id": intent.id, "name": intent.name} for intent in dataset.intents],
+            "samples": augmented_samples
+        }, f, indent=4, ensure_ascii=False)
+
+    splits = {}
+    max_num = max(args.split_numbers)
+    for n in args.split_numbers:
+        if n > max_num:
+            raise ValueError(f"Requested {n} examples but max is {max_num}")
+
+        class_to_samples = defaultdict(list)
+        for sample in augmented_samples:
+            class_to_samples[sample["label"]].append(sample)
+
+        selected = []
+        for class_id, samples in class_to_samples.items():
+            selected.extend(random.sample(samples, k=n))
+
+        splits[n] = selected
+
+    original_data = dataset["train"].to_list()
+    for n, aug_samples in splits.items():
+        combined = original_data + aug_samples
+
+        new_dataset = Dataset.from_dict({
+            "intents": dataset.intents,
+            "train": combined
+        })
+
+        output_path = Path(args.output_dir) / f"dataset_{n}_examples.json"
+        new_dataset.to_json(output_path)
+        print(f"Saved {len(combined)} examples to {output_path}")
+
+if __name__ == "__main__":
+    main()

scripts/data/construct_fewshot_multiclass_dataset.py

@@ -0,0 +1,48 @@
+import json
+import os
+from argparse import ArgumentParser
+from collections import defaultdict
+from random import seed, sample
+from autointent import Dataset
+
+def main() -> None:
+    parser = ArgumentParser(description="Create few-shot version of multiclass dataset")
+    parser.add_argument("--dataset-name", type=str, required=True, 
+                      help="Hugging Face dataset path (e.g. 'AutoIntent/massive_ru')")
+    parser.add_argument("--output-path", type=str, required=True,
+                      help="Path to save few-shot dataset")
+    parser.add_argument("--k-shots", type=int, required=True,
+                      help="Number of examples per class")
+    parser.add_argument("--split", type=str, default="train",
+                      help="Dataset split to process")
+    parser.add_argument("--seed", type=int, default=0,
+                      help="Random seed for reproducibility")
+    args = parser.parse_args()
+
+    seed(args.seed)
+
+    dataset = Dataset.from_hub(args.dataset_name)
+
+    class_to_examples = defaultdict(list)
+    for example in dataset[args.split]:
+        class_to_examples[example["label"]].append(example["utterance"])
+
+    fewshot_examples = []
+    for class_id, utterances in class_to_examples.items():
+        if len(utterances) < args.k_shots:
+            raise ValueError(f"Class {class_id} has only {len(utterances)} examples")
+
+        selected = sample(utterances, args.k_shots)
+        fewshot_examples.extend([
+            {"utterance": utt, "label": class_id} for utt in selected
+        ])
+
+    fewshot_dataset = Dataset.from_dict({
+        "intents": dataset.intents,
+        args.split: fewshot_examples
+    })
+
+    fewshot_dataset.to_json(args.output_path)
+
+if __name__ == "__main__":
+    main()