Working draft

src/fabricator/dataset_generator.py

@@ -61,7 +61,7 @@ def generate(
  num_samples_to_generate: int = 10,
  timeout_per_prompt: Optional[int] = None,
  log_every_n_api_calls: int = 25,
- dummy_response: Optional[Union[str, Callable]] = None
+ dummy_response: Optional[Union[str, Callable]] = None,
  ) -> Union[Dataset, Tuple[Dataset, Dataset]]:
  """Generate a dataset based on a prompt template and support examples.
  Optionally, unlabeled examples can be provided to annotate unlabeled data.
@@ -93,8 +93,11 @@ def generate(
  if fewshot_dataset:
  self._assert_fewshot_dataset_matches_prompt(prompt_template, fewshot_dataset)
 
- assert fewshot_sampling_strategy in [None, "uniform", "stratified"], \
- "Sampling strategy must be 'uniform' or 'stratified'"
+ assert fewshot_sampling_strategy in [
+ None,
+ "uniform",
+ "stratified",
+ ], "Sampling strategy must be 'uniform' or 'stratified'"
 
  if fewshot_dataset and not fewshot_sampling_column:
  fewshot_sampling_column = prompt_template.generate_data_for_column[0]
@@ -111,7 +114,7 @@ def generate(
  num_samples_to_generate,
  timeout_per_prompt,
  log_every_n_api_calls,
- dummy_response
+ dummy_response,
  )
 
  if return_unlabeled_dataset:
@@ -134,7 +137,6 @@ def _try_generate(
  """
 
  if dummy_response:
-
  if isinstance(dummy_response, str):
  logger.info(f"Returning dummy response: {dummy_response}")
  return dummy_response
@@ -152,7 +154,7 @@ def _try_generate(
  prediction = self.prompt_node.run(
  prompt_template=HaystackPromptTemplate(prompt=prompt_text),
  invocation_context=invocation_context,
- )[0]["results"]
+ )
  except Exception as error:
  logger.error(f"Error while generating example: {error}")
  return None
@@ -172,7 +174,7 @@ def _inner_generate_loop(
  num_samples_to_generate: int,
  timeout_per_prompt: Optional[int],
  log_every_n_api_calls: int = 25,
- dummy_response: Optional[Union[str, Callable]] = None
+ dummy_response: Optional[Union[str, Callable]] = None,
  ):
  current_tries_left = self._max_tries
  current_log_file = self._setup_log(prompt_template)
@@ -200,8 +202,11 @@ def _inner_generate_loop(
 
  if fewshot_dataset:
  prompt_labels, fewshot_examples = self._sample_fewshot_examples(
- prompt_template, fewshot_dataset, fewshot_sampling_strategy, fewshot_examples_per_class,
- fewshot_sampling_column
+ prompt_template,
+ fewshot_dataset,
+ fewshot_sampling_strategy,
+ fewshot_examples_per_class,
+ fewshot_sampling_column,
  )
 
  prompt_text = prompt_template.get_prompt_text(prompt_labels, fewshot_examples)
@@ -231,6 +236,7 @@ def _inner_generate_loop(
  f" {len(generated_dataset)} examples."
  )
  break
+ continue
 
  if len(prediction) == 1:
  prediction = prediction[0]
@@ -310,8 +316,9 @@ def _convert_prediction(self, prediction: str, target_type: type) -> Any:
  return target_type(prediction)
  except ValueError:
  logger.warning(
- "Could not convert prediction {} to type {}. "
- "Returning original prediction.", repr(prediction), target_type
+ "Could not convert prediction {} to type {}. " "Returning original prediction.",
+ repr(prediction),
+ target_type,
  )
  return prediction
 
@@ -321,21 +328,20 @@ def _sample_fewshot_examples(
  fewshot_dataset: Dataset,
  fewshot_sampling_strategy: str,
  fewshot_examples_per_class: int,
- fewshot_sampling_column: str
+ fewshot_sampling_column: str,
  ) -> Tuple[Union[List[str], str], Dataset]:
-
  if fewshot_sampling_strategy == "uniform":
  prompt_labels = choice(prompt_template.label_options, 1)[0]
- fewshot_examples = fewshot_dataset.filter(
- lambda example: example[fewshot_sampling_column] == prompt_labels
- ).shuffle().select(range(fewshot_examples_per_class))
+ fewshot_examples = (
+ fewshot_dataset.filter(lambda example: example[fewshot_sampling_column] == prompt_labels)
+ .shuffle()
+ .select(range(fewshot_examples_per_class))
+ )
 
  elif fewshot_sampling_strategy == "stratified":
  prompt_labels = prompt_template.label_options
  fewshot_examples = single_label_stratified_sample(
- fewshot_dataset,
- fewshot_sampling_column,
- fewshot_examples_per_class
+ fewshot_dataset, fewshot_sampling_column, fewshot_examples_per_class
  )
 
  else:
@@ -345,9 +351,11 @@ def _sample_fewshot_examples(
  else:
  fewshot_examples = fewshot_dataset.shuffle()
 
- assert len(fewshot_examples) > 0, f"Could not find any fewshot examples for label(s) {prompt_labels}." \
- f"Ensure that labels of fewshot examples match the label_options " \
- f"from the prompt."
+ assert len(fewshot_examples) > 0, (
+ f"Could not find any fewshot examples for label(s) {prompt_labels}."
+ f"Ensure that labels of fewshot examples match the label_options "
+ f"from the prompt."
+ )
 
  return prompt_labels, fewshot_examples
 

src/fabricator/nodes/__init__.py

@@ -0,0 +1,7 @@
+from .base import PromptNode
+from .nodes import GuidedPromptNode
+
+__all__ = [
+ "PromptNode",
+ "GuidedPromptNode",
+]

src/fabricator/nodes/base.py

@@ -0,0 +1,24 @@
+from abc import ABC, abstractmethod
+
+from typing import Optional, Dict, Union
+
+from haystack.nodes import PromptNode as HaystackPromptNode
+from haystack.nodes.prompt import PromptTemplate as HaystackPromptTemplate
+
+
+class Node(ABC):
+ @abstractmethod
+ def run(self, prompt_template):
+ pass
+
+
+class PromptNode(Node):
+ def __init__(self, model_name_or_path: str, *args, **kwargs) -> None:
+ self._prompt_node = HaystackPromptNode(model_name_or_path, *args, **kwargs)
+
+ def run(
+ self,
+ prompt_template: Optional[Union[str, HaystackPromptTemplate]],
+ invocation_context: Optional[Dict[str, any]] = None,
+ ):
+ return self._prompt_node.run(prompt_template, invocation_context=invocation_context)[0]["results"]

src/fabricator/nodes/nodes.py

@@ -0,0 +1,49 @@
+import json
+
+from typing import Optional, Dict, Union
+
+from haystack.nodes import PromptNode as HaystackPromptNode
+from haystack.nodes.prompt import PromptTemplate as HaystackPromptTemplate
+
+from .base import Node
+
+try:
+ import outlines.models as models
+ import outlines.text.generate as generate
+
+ from pydantic import BaseModel
+
+ import torch
+except ImportError as exc:
+ raise ImportError("Try 'pip install outlines'") from exc
+
+
+class GuidedPromptNode(Node):
+ def __init__(
+ self,
+ model_name_or_path: str,
+ schema: Union[str, BaseModel],
+ max_length: int = 100,
+ device: Optional[str] = None,
+ model_kwargs: Dict = None,
+ manual_seed: Optional[int] = None,
+ ) -> None:
+ self.max_length = max_length
+ model_kwargs = model_kwargs or {}
+ self._model = models.transformers(model_name_or_path, device=device, **model_kwargs)
+ # JSON schema of class
+ if not isinstance(schema, str):
+ schema = json.dumps(schema.schema())
+
+ self._generator = generate.json(
+ self._model,
+ schema,
+ self.max_length,
+ )
+
+ self.rng = torch.Generator(device=device)
+ if manual_seed is not None:
+ self.rng.manual_seed(manual_seed)
+
+ def run(self, prompt_template: HaystackPromptTemplate, **kwargs):
+ return self._generator(prompt_template.prompt_text, rng=self.rng)