QA dataset generation and rise experiment

.gitignore

@@ -209,10 +209,13 @@ studies/box*
 studies/*example1--small-study--drdocs_hf.yaml
 studies/silver*
 studies/test*
+studies/rise*
 
 notebooks/.nfs*
 
 /private
 *private*
 /flowgen/data/nltk-data
 .nfs*
+data.*
+/datasets

notebooks/create_dataset.ipynb

@@ -0,0 +1,364 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "0",
+   "metadata": {},
+   "source": [
+    "# QA Dataset Generation\n",
+    "Given a raw text, the notebook helps to generate a custom HuggingFace QA dataset based on the given information."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%reload_ext autoreload\n",
+    "%autoreload 2\n",
+    "\n",
+    "from IPython.core import ultratb\n",
+    "\n",
+    "ultratb.VerboseTB.tb_highlight = \"bg:#3e0054\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "if not os.getcwd().endswith(\"syftr\"):\n",
+    "    os.chdir(os.path.dirname(os.getcwd()))\n",
+    "    print(f\"Changed working directory to: {os.getcwd()}\")\n",
+    "\n",
+    "from syftr.configuration import cfg"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "DATA_FILEPATH = \"data.md\"  # Path to the raw text file\n",
+    "CHUNK_SIZE = 400  # Size of each text chunk\n",
+    "LLMS = [  # adjust to LLMs you want to use for question generation\n",
+    "    \"gpt-4o-mini\",\n",
+    "    \"Qwen/Qwen3-32B\",\n",
+    "    \"google/gemma-3-27b-it\",\n",
+    "]  # We randomly select one of the provided LLMs per chunk\n",
+    "NUM_PARALLEL = 50  # Number of parallel processes to use for chunk processing\n",
+    "# -------------------------------------------------------------------------------------------\n",
+    "  # Add instructions that are specific to your QA generation task\n",
+    "CUSTOM_QA_INSTRUCTIONS = None\n",
+    "\n",
+    "assert CUSTOM_QA_INSTRUCTIONS, \"Please provide custom instructions for the QA generation.\"\n",
+    "\n",
+    "# Provide a valid dataset name\n",
+    "DATASET_NAME = None\n",
+    "assert DATASET_NAME, \"Please set the DATASET_NAME variable to a valid dataset name.\"\n",
+    "# -------------------------------------------------------------------------------------------\n",
+    "\n",
+    "DATASET_IS_PRIVATE = True  # Set to False if you want to share the dataset publicly\n",
+    "\n",
+    "HF_DATASET_NAME = f\"DataRobot-Research/{DATASET_NAME}\"  # Adjust name of the dataset on Hugging Face Hub\n",
+    "HF_TOKEN = cfg.hf_datasets.api_key.get_secret_value()  # Get Hugging Face token from configuration\n",
+    "\n",
+    "assert HF_TOKEN, \"Please set the HF_TOKEN environment variable with your Hugging Face token.\"\n",
+    "\n",
+    "print(f\"Using Hugging Face token: {HF_TOKEN[:4]}...{HF_TOKEN[-4:]}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def load_text(file_path: str) -> str:\n",
+    "    with open(file_path, \"r\", encoding=\"utf-8\") as file:\n",
+    "        return file.read()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "raw_text = load_text(DATA_FILEPATH)\n",
+    "print(f\"Loaded {len(raw_text)} characters from {DATA_FILEPATH}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def chunk_text(text: str, chunk_size: int = 1000) -> list:\n",
+    "    return [text[i : i + chunk_size] for i in range(0, len(text), chunk_size)]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "chunks = chunk_text(raw_text, CHUNK_SIZE)\n",
+    "print(f\"Created {len(chunks)} chunks of size {CHUNK_SIZE} characters.\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from tenacity import retry, stop_after_attempt, wait_fixed\n",
+    "from syftr.llm import get_llm\n",
+    "\n",
+    "\n",
+    "@retry(stop=stop_after_attempt(5), wait=wait_fixed(2))\n",
+    "def generate(prompt: str, llm_name: str, **kwargs):\n",
+    "    llm = get_llm(llm_name)\n",
+    "    assert llm is not None, f\"LLM {llm_name} not found.\"\n",
+    "    response = llm.complete(prompt=prompt, **kwargs)\n",
+    "    return response.text"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def generate_qa_from_chunk(\n",
+    "    chunk: str, llm_name: str, **kwargs\n",
+    ") -> str:\n",
+    "    prompt = f\"\"\"Generate a question and answer based on the text below. Make sure to not use special formatting, like markdown, but formulate the question and the answer in a plan text format. Start with the question followed by the answer. The question should be clear and concise, and the answer should be informative and directly related to the question, for instance,\n",
+    "    \n",
+    "    Question: Who is in charge of the project SuperGold?\n",
+    "\n",
+    "    Answer: The project is led by Dr. Jane Smith.\n",
+    "\n",
+    "    Note that the question should always be specific, for instance, don't use generic terms like \"the text\" but always be specific about what you mean and use concrete names whereever possible. Same with images and tables: make sure you can specify which table or image your question is about or do not ask this question. The answer should be a direct response to the question, providing relevant information from the text chunk provided below.\n",
+    "    If you cannot generate a question and answer based on the text, return an empty string.\n",
+    "    Moreover, follow these custom instructions: \\n\\n{CUSTOM_QA_INSTRUCTIONS}\\n\\n\n",
+    "\n",
+    "    Chunk: \\n\\n{chunk}\"\"\"\n",
+    "    response = generate(prompt, llm_name, **kwargs)\n",
+    "    return response.strip()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import hashlib\n",
+    "import re\n",
+    "import typing as T\n",
+    "\n",
+    "def parse_qa_pairs(text: str, llm_name: str | None = None, chunk: str | None = None) -> T.List[T.Dict[str, str]]:  \n",
+    "    pattern = r\"Question:\\s*(.*?)\\s*Answer:\\s*(.*)\"\n",
+    "    matches = re.findall(pattern, text, re.DOTALL)\n",
+    "    parsed_pairs = []\n",
+    "    for question, answer in matches:\n",
+    "        pair = {\n",
+    "            \"id\": hashlib.md5(f\"{question.strip()}_{answer.strip()}\".encode()).hexdigest(),\n",
+    "            \"question\": question.strip(),\n",
+    "            \"answer\": answer.strip(),\n",
+    "        }\n",
+    "        if llm_name:\n",
+    "            pair[\"qtype\"] = llm_name\n",
+    "        if chunk:\n",
+    "            pair[\"gold_evidence\"] = [chunk.strip()]\n",
+    "        parsed_pairs.append(pair)\n",
+    "    return parsed_pairs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "11",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import random\n",
+    "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
+    "\n",
+    "\n",
+    "def get_qa_pairs_from_chunks(chunks: T.List[str]) -> T.List[T.Dict[str, str]]:\n",
+    "    qa_pairs = []\n",
+    "\n",
+    "    def _gen(chunk: str) -> T.List[T.Dict[str, str]]:\n",
+    "        llm_name = random.choice(LLMS)\n",
+    "        generated_text = generate_qa_from_chunk(chunk, llm_name, max_tokens=1024, temperature=0.7)\n",
+    "        return parse_qa_pairs(generated_text, llm_name, chunk)\n",
+    "\n",
+    "    with ThreadPoolExecutor(max_workers=NUM_PARALLEL) as executor:\n",
+    "        futures = [executor.submit(_gen, chunk) for chunk in chunks]\n",
+    "        results = [future.result() for future in futures]\n",
+    "        for pairs in results:\n",
+    "            qa_pairs.extend(pairs)\n",
+    "    return qa_pairs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "12",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "qa_pairs = get_qa_pairs_from_chunks(chunks)\n",
+    "print(f\"Generated {len(qa_pairs)} Q&A pairs from {len(chunks)} chunks.\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "13",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for pair in qa_pairs:\n",
+    "    print(\"-\"* 40)\n",
+    "    print(f\"LLM: {pair.get('qtype', 'Unknown')}\\nQuestion: {pair['question']}\\nAnswer: {pair['answer']}\\nEvidence: {pair.get('gold_evidence', 'N/A')}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "14",
+   "metadata": {},
+   "source": [
+    "**Adjust the parameters to make a custom split based on your needs and the amount of data generated.**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "15",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import datasets\n",
+    "\n",
+    "\n",
+    "def get_context(chunks: T.List[str]) -> str:\n",
+    "    full_context = \"\\n\".join(chunks)\n",
+    "    return full_context\n",
+    "\n",
+    "def prepare_hf_data(\n",
+    "        qa_pairs: T.List[T.Dict[str, str]], \n",
+    "        chunks: T.List[str] | None = None, \n",
+    "        all_grounding_data_for_each_partition = True,\n",
+    ") -> T.Tuple[datasets.DatasetDict, datasets.DatasetDict]:\n",
+    "    if all_grounding_data_for_each_partition:\n",
+    "        grounding_data_train = chunks\n",
+    "        grounding_data_test = chunks\n",
+    "        grounding_data_holdout = chunks\n",
+    "        grounding_data_sample = chunks[:5]\n",
+    "    elif chunks:\n",
+    "        grounding_data_train = get_context(chunks[:100])\n",
+    "        grounding_data_test = get_context(chunks[100:200])\n",
+    "        grounding_data_holdout = get_context(chunks[200:])\n",
+    "        grounding_data_sample = get_context(chunks[:5])\n",
+    "    else:\n",
+    "        raise ValueError(\"Either chunks or raw_text must be provided.\")\n",
+    "    \n",
+    "    qa_data = datasets.DatasetDict(\n",
+    "        {\n",
+    "            \"train\": datasets.Dataset.from_list(qa_pairs[:100]),\n",
+    "            \"test\": datasets.Dataset.from_list(qa_pairs[100:200]),\n",
+    "            \"holdout\": datasets.Dataset.from_list(qa_pairs[200:]),\n",
+    "            \"sample\": datasets.Dataset.from_list(qa_pairs[:5]),  # for quick testing\n",
+    "        }\n",
+    "    )\n",
+    "    grounding_data = datasets.DatasetDict(\n",
+    "        {\n",
+    "            \"train\": datasets.Dataset.from_dict({\"text\": grounding_data_train}),\n",
+    "            \"test\": datasets.Dataset.from_dict({\"text\": grounding_data_test}),\n",
+    "            \"holdout\": datasets.Dataset.from_dict({\"text\": grounding_data_holdout}),\n",
+    "            \"sample\": datasets.Dataset.from_dict({\"text\": grounding_data_sample}),\n",
+    "        }\n",
+    "    )\n",
+    "    return qa_data, grounding_data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "16",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "qa_data, grounding_data = prepare_hf_data(qa_pairs, chunks=chunks)\n",
+    "\n",
+    "qa_data.push_to_hub(\n",
+    "    repo_id=HF_DATASET_NAME, \n",
+    "    data_dir=\"examples\",\n",
+    "    private=DATASET_IS_PRIVATE, \n",
+    "    token=HF_TOKEN,\n",
+    "    config_name=\"qa\"\n",
+    ")\n",
+    "print(f\"QA data pushed to Hugging Face Hub.\")\n",
+    "\n",
+    "grounding_data.push_to_hub(\n",
+    "    repo_id=HF_DATASET_NAME,\n",
+    "    data_dir=\"grounding_data\",\n",
+    "    private=DATASET_IS_PRIVATE,\n",
+    "    token=HF_TOKEN,\n",
+    "    config_name=\"grounding\"\n",
+    ")\n",
+    "print(f\"Grounding data pushed to Hugging Face Hub.\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "17",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "syftr",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}