cosmosage.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# cosmosage\n",
    "\n",
    "See README for more information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import os\n",
    "import scrape_arxiv\n",
    "import analyze_asl_dict\n",
    "import glob\n",
    "import random\n",
    "import multiprocessing\n",
    "import json\n",
    "import time\n",
    "from langchain.chains.summarize import load_summarize_chain\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "import glob\n",
    "import random"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Choose arXiv papers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cache_file = \"datasets/arxiv_ids_cache.pkl\"\n",
    "\n",
    "# Check if the cache file exists\n",
    "if os.path.exists(cache_file):\n",
    "    # Load the cached data\n",
    "    with open(cache_file, \"rb\") as f:\n",
    "        arxiv_ids = pickle.load(f)\n",
    "else:\n",
    "    # unique arXiv numbers from the asl database\n",
    "    db_path = \"datasets/dict_20231123.db\"\n",
    "    arxiv_id_asl_tagged = analyze_asl_dict.extract_unique_arxiv_numbers(db_path)\n",
    "\n",
    "    # also extract all of my papers\n",
    "    search_params = {\"search_query\": \"au:de_Haan_T\", \"searchtype\": \"author\"}\n",
    "    arxiv_id_tdh = scrape_arxiv.get_arxiv_ids(search_params)\n",
    "\n",
    "    # also extract the papers with \"cosmic microwave background\" in the abstract\n",
    "    search_params = {\"search_query\": \"abs:\\\"cosmic microwave background\\\"\"}\n",
    "    arxiv_id_cmb = scrape_arxiv.get_arxiv_ids(search_params)\n",
    "\n",
    "    # more arxiv papers recommended for me by asl\n",
    "    arxiv_id_asl_rec = scrape_arxiv.other_arxiv_recommendation_ids()\n",
    "\n",
    "    # join all of these arxiv ids and remove duplicates\n",
    "    arxiv_ids = arxiv_id_asl_tagged + arxiv_id_tdh + arxiv_id_cmb + arxiv_id_asl_rec\n",
    "    arxiv_ids = list(set(arxiv_ids))\n",
    "\n",
    "    # Save the data to the cache file\n",
    "    with open(cache_file, \"wb\") as f:\n",
    "        pickle.dump(arxiv_ids, f)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Synthetic data generation on arXiv papers\n",
    "\n",
    "Here, we generate synthetic data using the following:\n",
    " - instruction-tuned model to generate the QA pairs & summaries\n",
    " - VLLM server to load the model once and provide good throughput\n",
    " - langchain to handle \n",
    "   - gathering of papers\n",
    "   - extracting from PDFs\n",
    "   - chunking data\n",
    "   - summarization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2.1: Download, parse, and cache the arXiv papers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# filter out any arxiv ids that don't start with 07 through 24 since the process will fail on these anyway\n",
    "new_arxiv_ids = []\n",
    "for arxiv_id in arxiv_ids:\n",
    "    outfile = f\"datasets/arxiv_cache/{arxiv_id}.pkl\"\n",
    "    if os.path.exists(outfile):\n",
    "        continue\n",
    "    id_year = int(arxiv_id[0:2])\n",
    "    if id_year >= 7 and id_year <= 24:\n",
    "        new_arxiv_ids.append(arxiv_id)\n",
    "print(f\"Found {len(new_arxiv_ids)} arxiv ids from 2007 to 2024.\")\n",
    "\n",
    "print(\"Shuffling arxiv ids...\")\n",
    "random.shuffle(new_arxiv_ids)\n",
    "\n",
    "def cache_arxiv_id(arxiv_id):\n",
    "    outfile = f\"datasets/arxiv_cache/{arxiv_id}.pkl\"\n",
    "    if os.path.exists(outfile):\n",
    "        print(f\"File {outfile} already exists, skipping.\")\n",
    "        return\n",
    "    print(f\"Processing {arxiv_id}\\n\")\n",
    "    try:\n",
    "        print(f\"Waiting 2 seconds before scraping {arxiv_id}...\")\n",
    "        time.sleep(2)\n",
    "        paper = scrape_arxiv.ArxivPaper(arxiv_id)\n",
    "        paper.save_to_cache()\n",
    "    except Exception as e:\n",
    "        print(f\"Error processing {arxiv_id}: {e}\\n\")\n",
    "\n",
    "# Create a pool of workers\n",
    "pool = multiprocessing.Pool(1)\n",
    "\n",
    "# Map the process_arxiv_id function to each arxiv_id in parallel\n",
    "for arxiv_id in new_arxiv_ids:\n",
    "    pool.apply_async(cache_arxiv_id, args=(arxiv_id,))\n",
    "\n",
    "# Close the pool of workers\n",
    "pool.close()\n",
    "pool.join()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2.2 Add summaries\n",
    "\n",
    "For the generation of summaries, you can use the following code block. However, to speed things up, you can also run the script `generate_summaries_standalone.py` which can be run in parallel on e.g. a GPU cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add summaries to the pkl files\n",
    "filenames = glob.glob(\"datasets/arxiv_cache/*.pkl\")\n",
    "random.shuffle(filenames)\n",
    "\n",
    "def add_summary(filename):\n",
    "    with open(filename, \"rb\") as f:\n",
    "        print(f\"Loading {filename}...\")\n",
    "        paper = pickle.load(f)\n",
    "        if \"summary\" in paper:\n",
    "            print(f\"Skipping {filename}, already has summary.\")\n",
    "            return\n",
    "\n",
    "        # make sure to start a vLLM ChatOpenAI server first\n",
    "        inference_server_url = \"http://0.0.0.0:8000/v1\"\n",
    "\n",
    "        llm = ChatOpenAI(\n",
    "            model=\"/home/tijmen/cosmosage/packages/text-generation-webui/models/TheBloke_bagel-dpo-34b-v0.2-GPTQ_gptq-4bit-32g-actorder_True\",\n",
    "            openai_api_key=\"EMPTY\",\n",
    "            openai_api_base=inference_server_url,\n",
    "            temperature=0.4,\n",
    "        )\n",
    "\n",
    "        print(\"Generating summary...\")\n",
    "        summarize_chain = load_summarize_chain(llm, chain_type=\"map_reduce\")\n",
    "        summary = summarize_chain.run(paper.pages)\n",
    "        print(f\"Summary generated. Length {len(summary)}\")\n",
    "\n",
    "        paper[\"summary\"] = summary\n",
    "\n",
    "    if \"summary\" in paper:\n",
    "        if len(summary)>10:\n",
    "            with open(filename, \"wb\") as f:\n",
    "                print(\"Saving summary...\")\n",
    "                pickle.dump(paper, f)\n",
    "\n",
    "# Create a pool of workers\n",
    "pool = multiprocessing.Pool(1)\n",
    "\n",
    "# Map the process_arxiv_id function to each arxiv_id in parallel\n",
    "for filename in filenames:\n",
    "    pool.apply_async(add_summary, args=(filename,))\n",
    "\n",
    "# Close the pool of workers\n",
    "pool.close()\n",
    "pool.join()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2.3 Generate QA pairs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For generation of QA pairs, run `generate_synth_standalone.py`. This script can be run many times in parallel to speed up the process."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Synthetic data generation on textbooks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "from extract_textbooks import TextBook\n",
    "textbooks = []\n",
    "for filepath in glob.glob(\"datasets/cosmology_textbooks/*.txt\"):\n",
    "    textbooks.append(TextBook(filepath))\n",
    "for textbook in textbooks:\n",
    "    textbook.generate_qa_pairs(multiprocess=True)\n",
    "    textbook.save_dataset_jsonl()\n",
    "    print(f\"Saved {textbook.author} to jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "# collate all the JSONL files and shuffle them for good measure\n",
    "textbook_jsonl_files_in = glob.glob(\"datasets/cosmology_textbooks_qa/*/*.jsonl\")\n",
    "textbook_jsonl_file_out = \"datasets/cosmology_textbooks_qa.jsonl\"\n",
    "\n",
    "with open(textbook_jsonl_file_out, \"w\") as f:\n",
    "    all_lines = []\n",
    "    for textbook_jsonl_file in textbook_jsonl_files_in:\n",
    "        with open(textbook_jsonl_file, \"r\") as g:\n",
    "            all_lines.extend(g.readlines())\n",
    "    random.shuffle(all_lines)\n",
    "    f.writelines(all_lines)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Join, prepare the datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Depending on the choices we made in Steps 2 and 3, we will now have a bunch of synthetic summaries, QA pairs, and other data from arXiv papers and textbooks. We may also have other sources such as public datasets. \n",
    "\n",
    "The goal of Step 4 is to get all these datasets in JSONL format to get them ready for training with `fine_tune.py`, `fine_tune_lora.py`, or `axolotl`. We will go through and prepare any datasets that aren't yet in a useable JSONL format. One recurring theme will be the use of randomized choices in order to increase training set diversity. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4.1 Make training data from the summaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert summaries to QA format\n",
    "filenames = glob.glob(\"datasets/arxiv_cache/*.pkl\")\n",
    "random.shuffle(filenames)\n",
    "outfile = \"datasets/arxiv_summary3.jsonl\"\n",
    "with open(outfile, \"w\") as fout:\n",
    "    for filename in filenames:\n",
    "        with open(filename, \"rb\") as f:\n",
    "            paper = pickle.load(f)\n",
    "            if \"summary\" not in paper:\n",
    "                print(f\"Warning: no summary found for {paper['arxiv_id']} in {filename}.\")\n",
    "                continue\n",
    "            system_prompts = [\n",
    "                \"You are an expert cosmologist. You provide answers to questions about one particular paper.\",\n",
    "                \"You are a knowledgeable cosmologist aware of the latest research. Follow the user's request.\",\n",
    "                \"You are a seasoned astrophysicist. Provide insights into an cosmology or cosmology paper.\",\n",
    "                \"As an expert in the field of cosmology, offer an explanation of whatever paper the user is asking about.\",\n",
    "                \"You have a deep understanding of astrophysical research. Guide the user through the main points of whatever study they are asking about.\",\n",
    "                \"You are adept in explaining complex astrophysical concepts. Explain about the user's requested topic.\",\n",
    "                \"As a cosmology specialist, distill the essence of whatever the user is asking about.\"\n",
    "            ]\n",
    "            system_prompt = random.choice(system_prompts)\n",
    "            user_prompts = [\n",
    "                f\"Please summarize {paper['shorthand_title']}.\",\n",
    "                f\"What is {paper['arxiv_id']} about?\",\n",
    "                f\"Can you explain the key findings of {paper['shorthand_title']}?\",\n",
    "                f\"What are the major contributions of {paper['shorthand_title']} to cosmology?\",\n",
    "                f\"I'm interested in the contents of the paper {paper['arxiv_id']}. Can you describe it?\",\n",
    "                f\"What's {paper['shorthand_title']}?\",\n",
    "                f\"Please explain the main points of {paper['shorthand_title']}.\",\n",
    "            ]\n",
    "            user_prompt = random.choice(user_prompts)\n",
    "            assistant_preambles = [\n",
    "                f\"{paper['shorthand_title']} is titled \\\"{paper['title']}\\\". Here is a summary of the paper. \",\n",
    "                \"\", # no preamble\n",
    "                f\"In the paper titled \\\"{paper['title']}\\\", published in {paper['year']}, the authors explore an intriguing aspect of cosmology. Let's delve into the summary. \",\n",
    "                f\"The study \\\"{paper['title']}\\\" provides key insights. Here's an overview. \",\n",
    "                f\"Delving into \\\"{paper['title']}\\\", the summary goes as follows. \",\n",
    "                f\"This summary focuses on \\\"{paper['title']}\\\", a noteworthy paper with the arXiv ID {paper['arxiv_id']}. The key points are the following. \",\n",
    "                f\"Titled \\\"{paper['title']}\\\", this paper presents groundbreaking research in cosmology. Here's a summary. \",\n",
    "                f\"\\\"{paper['title']}\\\" is a paper about cosmology. It discusses the following. \",\n",
    "                \"Summary. \",\n",
    "            ]\n",
    "            assistant_preamble = random.choice(assistant_preambles)\n",
    "            assistant_message = assistant_preamble+paper['summary']\n",
    "            # sharegpt should take the form {\"conversations\": [{\"from\": \"...\", \"value\": \"...\"}]} where the message is from \"system\", \"human\", or \"gpt\".\n",
    "            conversation_data = {\n",
    "                \"conversations\": [\n",
    "                    {\"from\": \"system\", \"value\": system_prompt},\n",
    "                    {\"from\": \"human\", \"value\": user_prompt},\n",
    "                    {\"from\": \"gpt\", \"value\": assistant_message}\n",
    "                ]\n",
    "            }\n",
    "\n",
    "            fout.write(json.dumps(conversation_data) + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4.2 Make training data just from the metadata alone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filenames = glob.glob(\"datasets/arxiv_cache/*.pkl\")\n",
    "random.shuffle(filenames)\n",
    "outfile = \"datasets/arxiv_metadata_qa3.jsonl\"\n",
    "n_pairs_per_paper = 5\n",
    "system_prompts = [\n",
    "    \"You are an AI programmed to provide brief, factual answers about arXiv papers.\",\n",
    "    \"Your responses should be concise and limited to the essential details from the arXiv database.\",\n",
    "    \"Provide short, precise answers with just the key information from the arXiv papers.\",\n",
    "    \"You are an efficient AI, capable of giving terse answers about specific arXiv papers.\",\n",
    "    \"Deliver quick and factual responses about arXiv papers, focusing only on the core details.\",\n",
    "    \"Your role is to provide succinct, accurate information from the arXiv papers in as few words as possible.\",\n",
    "    \"As an AI, offer brief and direct answers about arXiv papers, omitting any extraneous details.\",\n",
    "    \"You are configured to give short, to-the-point responses about arXiv paper details.\",\n",
    "    \"Your task is to provide the most important information from the arXiv papers in a concise manner.\",\n",
    "    \"Respond with only the essential facts from the arXiv papers, keeping your answers brief.\",\n",
    "    \"You are an AI trained to deliver short and accurate summaries of information from arXiv papers.\",\n",
    "    \"Provide compact, factual responses, focusing solely on the critical aspects of arXiv papers.\",\n",
    "    \"You are an AI assistant. Provide only brief, metadata-based responses, without explaining the content.\"\n",
    "]\n",
    "\n",
    "qa_templates = [\n",
    "{\n",
    "    \"question\": \"Where can I find the arXiv paper titled {full_title} by {first_author} et al.?\",\n",
    "    \"answer\": \"The paper titled '{full_title}' by {first_author} et al. can be found at http://arxiv.org/abs/{arxiv_id}\"\n",
    "},\n",
    "{\n",
    "    \"question\": \"What is the arXiv ID for {first_author} et al.?\",\n",
    "    \"answer\": \"The arXiv ID for '{full_title}' by {first_author} et al. is {arxiv_id}.\"\n",
    "},\n",
    "{\n",
    "    \"question\": \"Who is the first author of the paper with arXiv ID {arxiv_id}?\",\n",
    "    \"answer\": \"The first author of {full_title} is {first_author}.\"\n",
    "},\n",
    "{\n",
    "    \"question\": \"What year was the arXiv paper {full_title} by {first_author} et al. published?\",\n",
    "    \"answer\": \"{first_author} et al. was published in {year}.\"\n",
    "},\n",
    "{\n",
    "    \"question\": \"Can you give me the link to the arXiv paper authored by {first_author} et al. in {year}?\",\n",
    "    \"answer\": \"The paper authored by {first_author} et al. in {year} can be found at http://arxiv.org/abs/{arxiv_id}.\"\n",
    "},\n",
    "{\n",
    "    \"question\": \"What is the topic of the paper by {first_author} et al. with arXiv ID {arxiv_id}?\",\n",
    "    \"answer\": \"{first_author} et al. discuss '{full_title}' in their paper with arXiv ID {arxiv_id}.\"\n",
    "},\n",
    "{\n",
    "    \"question\": \"How can I access {first_author} et al.'s {year} paper on the arXiv?\",\n",
    "    \"answer\": \"You can access the {year} paper by {first_author} et al. on arXiv at http://arxiv.org/abs/{arxiv_id}.\"\n",
    "},\n",
    "{\n",
    "    \"question\": \"Is the paper titled '{full_title}' available on arXiv?\",\n",
    "    \"answer\": \"Yes, the paper titled '{full_title}' is available on arXiv under ID {arxiv_id}.\"\n",
    "},\n",
    "{\n",
    "    \"question\": \"What research did {first_author} et al. present in {year} on arXiv?\",\n",
    "    \"answer\": \"In {year}, {first_author} et al. presented research on '{full_title}', available on arXiv with ID {arxiv_id}.\"\n",
    "},\n",
    "{\n",
    "    \"question\": \"Where can I read about {first_author} et al.'s findings in '{full_title}'?\",\n",
    "    \"answer\": \"You can read about {first_author} et al.'s findings in '{full_title}' at http://arxiv.org/abs/{arxiv_id} on arXiv.\"\n",
    "},\n",
    "{\n",
    "    \"question\": \"Can you provide a link to {first_author} et al.'s work titled '{full_title}' on arXiv?\",\n",
    "    \"answer\": \"Sure, the link to '{full_title}' by {first_author} et al. on arXiv is http://arxiv.org/abs/{arxiv_id}.\"\n",
    "},\n",
    "{\n",
    "    \"question\": \"What is the subject of the arXiv paper with ID {arxiv_id} by {first_author} et al.?\",\n",
    "    \"answer\": \"The subject of the arXiv paper with ID {arxiv_id} by {first_author} et al. is '{full_title}'.\"\n",
    "},\n",
    "{\n",
    "    \"question\": \"Who led the research for the paper on arXiv titled '{full_title}'?\",\n",
    "    \"answer\": \"{first_author} led the research for the paper titled '{full_title}' on arXiv.\"\n",
    "},\n",
    "{\n",
    "    \"question\": \"When was the paper titled '{full_title}' added to arXiv?\",\n",
    "    \"answer\": \"The paper titled '{full_title}' was added to arXiv in {year}.\"\n",
    "},\n",
    "{\n",
    "    \"question\": \"I need a copy of '{full_title}' by {first_author} et al. Can you help find it?\",\n",
    "    \"answer\": \"Here's the link where you can get your hands on '{full_title}' by {first_author} et al.: http://arxiv.org/abs/{arxiv_id}\"\n",
    "},\n",
    "]\n",
    "with open(outfile, \"w\") as fout:\n",
    "    for filename in filenames:\n",
    "        with open(filename, \"rb\") as f:\n",
    "            paper = pickle.load(f)\n",
    "            for _ in range(n_pairs_per_paper):\n",
    "                system_prompt = random.choice(system_prompts)\n",
    "                qa_template = random.choice(qa_templates)\n",
    "                question = qa_template[\"question\"].format(\n",
    "                    first_author=paper[\"first_author\"],\n",
    "                    year=paper[\"year\"],\n",
    "                    full_title=paper[\"title\"],\n",
    "                    arxiv_id=paper[\"arxiv_id\"],\n",
    "                    shorthand_title=paper[\"shorthand_title\"],\n",
    "                )\n",
    "                answer = qa_template[\"answer\"].format(\n",
    "                    first_author=paper[\"first_author\"],\n",
    "                    year=paper[\"year\"],\n",
    "                    full_title=paper[\"title\"],\n",
    "                    arxiv_id=paper[\"arxiv_id\"],\n",
    "                    shorthand_title=paper[\"shorthand_title\"],\n",
    "                )\n",
    "                # sharegpt should take the form {\"conversations\": [{\"from\": \"...\", \"value\": \"...\"}]} where the message is from \"system\", \"human\", or \"gpt\".\n",
    "                conversation_data = {\n",
    "                    \"conversations\": [\n",
    "                        {\"from\": \"system\", \"value\": system_prompt},\n",
    "                        {\"from\": \"human\", \"value\": question},\n",
    "                        {\"from\": \"gpt\", \"value\": answer}\n",
    "                    ]\n",
    "                }\n",
    "\n",
    "                fout.write(json.dumps(conversation_data) + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4.3 Make training data from the generated QA pairs.\n",
    "\n",
    "The QA information was generated by letting the model read chunks of text and write QA pairs. Then in a second pass, the model was told to critique, grade, write an alternative, grade. Let's start by analyze these grades."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import random\n",
    "import pickle\n",
    "papers = []\n",
    "filenames = glob.glob(\"datasets/arxiv_cache/*.pkl\")\n",
    "random.shuffle(filenames)\n",
    "outfile = \"datasets/arxiv_qa3.jsonl\"\n",
    "with open(outfile, \"w\") as fout:\n",
    "    for filename in filenames:\n",
    "        with open(filename, \"rb\") as f:\n",
    "            paper = pickle.load(f)\n",
    "            if \"qa\" not in paper:\n",
    "                print(f\"Warning: no QA pairs found for {paper['arxiv_id']} in {filename}.\")\n",
    "                continue\n",
    "            if len(paper[\"qa\"]) < 1:\n",
    "                print(f\"Warning: empty QA pairs found for {paper['arxiv_id']} in {filename}.\")\n",
    "                continue\n",
    "            if \"refined_answer\" not in paper[\"qa\"][0]:\n",
    "                print(f\"Warning: no refined answers found for {paper['arxiv_id']} in {filename}.\")\n",
    "                continue\n",
    "            papers.append(paper)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "student_grades = []\n",
    "teacher_grades = []\n",
    "for paper in papers:\n",
    "    for qa in paper['qa']:\n",
    "        student_grades.append(qa['student_grade'])\n",
    "        teacher_grades.append(qa['teacher_grade'])\n",
    "\n",
    "def grade_to_int(grade):\n",
    "    if isinstance(grade, int):\n",
    "        if 0 <= grade <= 100:\n",
    "            return grade\n",
    "        else:\n",
    "            print(f\"Error: grade '{grade}' is an int, but not in the range 0-100.\")\n",
    "            return 0\n",
    "    if isinstance(grade, dict):\n",
    "        for key in ['student', 'teacher', 'score', 'grade', 'value', 'complete_answer', 'total', 'percentile', 'int', 'raw', 'overall', 'explanation', 'age']:\n",
    "            if key in grade:\n",
    "                return grade_to_int(grade[key])  # Recursively process the extracted value\n",
    "    elif isinstance(grade, list) and grade:\n",
    "        return grade_to_int(grade[0])  # Recursively process the first item\n",
    "    elif isinstance(grade, str):\n",
    "        if grade.startswith('grade '):\n",
    "            grade = grade[6:]  # Skip 'grade '\n",
    "        try:\n",
    "            return int(grade)\n",
    "        except ValueError:\n",
    "            print(f\"Cannot convert string '{grade}' to int.\")\n",
    "            return 0\n",
    "    else:\n",
    "        print(f\"Error: grade '{grade}' is a {type(grade)}, which cannot be converted to int.\")\n",
    "        return 0\n",
    "    # Handle any other unanticipated cases\n",
    "    return 0  # Default return for unhandled cases\n",
    "\n",
    "# Convert student and teacher grades to integers\n",
    "import numpy as np\n",
    "student_grades_arr = np.array([grade_to_int(grade) for grade in student_grades])\n",
    "teacher_grades_arr = np.array([grade_to_int(grade) for grade in teacher_grades])\n",
    "# zero out any None\n",
    "student_grades_arr = np.nan_to_num(student_grades_arr)\n",
    "teacher_grades_arr = np.nan_to_num(teacher_grades_arr)\n",
    "print(f\"The average student grade is {np.mean(student_grades_arr):.2f} with a standard deviation of {np.std(student_grades_arr):.2f}.\")\n",
    "print(f\"The average teacher grade is {np.mean(teacher_grades_arr):.2f} with a standard deviation of {np.std(teacher_grades_arr):.2f}.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot separate histograms for student and teacher grades\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.hist(student_grades_arr, bins=np.linspace(0, 100, 100+1), alpha=0.5, label='Student')\n",
    "plt.hist(teacher_grades_arr, bins=np.linspace(0, 100, 100+1), alpha=0.5, label='Teacher')\n",
    "\n",
    "plt.xlabel(\"Grade\")\n",
    "plt.ylabel(\"Count\")\n",
    "plt.title(\"Distribution of Student and Teacher Grades\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From this, let's choose the following rules:\n",
    " - if the student gets a 90% or above, use their answer\n",
    " - else if the teacher gets 90% of above use their answer\n",
    " - else disregard this QA pair"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import random\n",
    "import json\n",
    "conversations = []\n",
    "for paper in papers:\n",
    "    system_prompts_short = [\n",
    "        \"You are an AI programmed to provide brief answers about arXiv papers.\",\n",
    "        \"As an expert cosmologist, you provide concise answers to the user's questions about an arXiv paper.\",\n",
    "        \"You are a specialized AI, offering succinct insights into cosmology research papers.\",\n",
    "        \"Provide brief, yet informative insights from arXiv papers in cosmology.\",\n",
    "        \"As a focused expert in cosmology, respond with short, precise answers about scientific papers.\",\n",
    "        \"Deliver quick and concise explanations of complex concepts from recent cosmology papers.\",\n",
    "        \"You are programmed to give short, clear responses to queries about arXiv cosmology publications.\",\n",
    "        \"Offer concise and direct answers to technical questions on cosmology research.\",\n",
    "        \"As an AI trained in cosmology, provide succinct responses to detailed scientific inquiries.\",\n",
    "        \"You are an AI designed to give brief, accurate descriptions of arXiv cosmology papers.\",\n",
    "        \"Quickly decipher and explain the key points from cosmology papers.\",\n",
    "        \"As an expert in the field, deliver short and precise interpretations of cosmology research.\",\n",
    "        \"Condense long cosmology papers into brief, understandable answers.\",\n",
    "        \"You are a compact knowledge source for quick insights into cosmology papers.\",\n",
    "        \"Provide to-the-point, accurate summaries of recent findings in cosmology.\",\n",
    "        \"As an AI, offer concise, clear-cut answers about specific aspects of cosmology papers.\",\n",
    "        \"Distill complex cosmology concepts from research papers into brief explanations.\",\n",
    "        \"Deliver quick, expert responses to questions about detailed cosmology research.\",\n",
    "        \"You are programmed to succinctly answer queries on advanced cosmology topics.\",\n",
    "        \"Offer crisp, clear summaries of the key findings in new cosmology research papers.\",\n",
    "    ]\n",
    "    system_prompt_medium = [\n",
    "        \"In a few sentences, explain the main findings of arXiv papers in cosmology.\",\n",
    "        \"Provide answers to the user's cosmology questions. Stick to a moderately long answer.\",\n",
    "        \"Elaborate briefly on the methodologies and conclusions of recent cosmology papers from arXiv.\",\n",
    "        \"Provide detailed, yet concise explanations of key theories and discoveries in recent cosmology papers.\",\n",
    "        \"As a cosmology expert, deliver medium-length answers that clarify complex concepts in arXiv papers.\",\n",
    "        \"Summarize the critical points of arXiv cosmology papers in a clear, moderately detailed manner.\",\n",
    "        \"Explain the significance and implications of findings in recent cosmology research, in a few sentences.\",\n",
    "        \"Interpret and convey the essence of cosmology papers from arXiv, aiming for moderate-length responses.\",\n",
    "        \"In a moderate amount of detail, discuss the innovations and findings in contemporary cosmology papers.\",\n",
    "        \"Delve into the core aspects of arXiv cosmology papers, providing answers that are informative yet succinct.\",\n",
    "        \"Shed light on the complexities of cosmology research with moderately expansive answers.\",\n",
    "        \"Analyze and explain the key aspects of arXiv cosmology papers in a clear, moderately lengthy format.\",\n",
    "    ]\n",
    "    system_prompt_long = [\n",
    "        \"Provide detailed answers to the user's questions about arXiv papers in cosmology.\",\n",
    "        \"Explain the methodologies and conclusions of recent cosmology papers from arXiv.\",\n",
    "        \"Provide detailed explanations of key theories and discoveries in recent cosmology papers.\",\n",
    "        \"As a cosmology expert, deliver long answers that clarify complex concepts in arXiv papers.\",\n",
    "        \"Explain the significance and implications of findings in recent cosmology research.\",\n",
    "        \"Interpret and convey the essence of cosmology papers from arXiv, aiming for longer responses.\",\n",
    "        \"In a detailed manner, discuss the innovations and findings in contemporary cosmology papers.\",\n",
    "        \"Delve into the core aspects of arXiv cosmology papers, providing answers that are informative yet detailed.\",\n",
    "        \"Shed light on the complexities of cosmology research with long answers.\",\n",
    "        \"Analyze and explain the key aspects of arXiv cosmology papers in a clear, lengthy format.\",\n",
    "    ]\n",
    "    for qa in paper[\"qa\"]:\n",
    "        question = qa[\"question\"]\n",
    "        if grade_to_int(qa[\"student_grade\"]) >= 90:\n",
    "            answer = qa[\"answer\"]\n",
    "        elif grade_to_int(qa[\"teacher_grade\"]) >= 90:\n",
    "            answer = qa[\"refined_answer\"]\n",
    "        else:\n",
    "            continue\n",
    "        qa_length = len(qa[\"answer\"])\n",
    "        if qa_length < 40:\n",
    "            continue\n",
    "        elif qa_length < 200:\n",
    "            system_prompt = random.choice(system_prompts_short)\n",
    "        elif qa_length < 500:\n",
    "            system_prompt = random.choice(system_prompt_medium)\n",
    "        else:\n",
    "            system_prompt = random.choice(system_prompt_long)\n",
    "\n",
    "        def replace_vague_phrases(text, replacements):\n",
    "            for phrase, replacement in replacements.items():\n",
    "                text = re.sub(phrase, replacement, text)\n",
    "            return text\n",
    "\n",
    "        # Dictionary of replacements. This will be iterated over, so the order matters. This is safe in python>=3.7\n",
    "        replacements = {\n",
    "            \"the analysis\": paper[\"shorthand_title\"],\n",
    "            \"The analysis\": paper[\"shorthand_title\"],\n",
    "            \"The passage\": paper[\"shorthand_title\"],\n",
    "            \"the passage\": paper[\"shorthand_title\"],\n",
    "            \"the PASSAGE\": paper[\"shorthand_title\"],\n",
    "            \"The PASSAGE\": paper[\"shorthand_title\"],\n",
    "            \"the paper\": paper[\"shorthand_title\"],\n",
    "            \"The paper\": paper[\"shorthand_title\"],\n",
    "            \"the study\": paper[\"shorthand_title\"],\n",
    "            \"The study\": paper[\"shorthand_title\"],\n",
    "            \"the research\": paper[\"shorthand_title\"],\n",
    "            \"The research\": paper[\"shorthand_title\"],\n",
    "            \"the authors' findings\": paper['shorthand_title'],\n",
    "            \"The authors' findings\": paper['shorthand_title'],\n",
    "            \"the authors\": f\"{paper['first_author']} et al.\",\n",
    "            \"The authors\": f\"{paper['first_author']} et al.\",\n",
    "            \"the findings\": paper[\"shorthand_title\"],\n",
    "            \"The findings\": paper[\"shorthand_title\"],\n",
    "            \"the method\": f\"the method used in {paper['shorthand_title']}\",\n",
    "            \"The method\": f\"The method used in {paper['shorthand_title']}\",\n",
    "        }\n",
    "\n",
    "        # Apply replacements to user prompt and assistant message\n",
    "        user_prompt = replace_vague_phrases(qa[\"question\"], replacements)\n",
    "        assistant_message = replace_vague_phrases(qa[\"answer\"], replacements)\n",
    "\n",
    "        # sharegpt should take the form {\"conversations\": [{\"from\": \"...\", \"value\": \"...\"}]} where the message is from \"system\", \"human\", or \"gpt\".\n",
    "        conversation_data = {\n",
    "            \"conversations\": [\n",
    "                {\"from\": \"system\", \"value\": system_prompt},\n",
    "                {\"from\": \"human\", \"value\": user_prompt},\n",
    "                {\"from\": \"gpt\", \"value\": assistant_message}\n",
    "            ]\n",
    "        }\n",
    "\n",
    "        conversations.append(json.dumps(conversation_data))\n",
    "\n",
    "random.shuffle(conversations)\n",
    "\n",
    "outfile = \"datasets/qa_tune/arxiv_refined_qa.jsonl\"\n",
    "with open(outfile, \"w\") as fout:\n",
    "    for conversation in conversations:\n",
    "        fout.write(conversation + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4.4 Add miscellanous datasets\n",
    "\n",
    "Here you can add whatever other datasets you might like."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4.5 Combine and shuffle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import random\n",
    "\n",
    "def read_jsonl(file_path):\n",
    "    \"\"\"Read a JSON Lines file and return a list of JSON objects.\"\"\"\n",
    "    with open(file_path, 'r', encoding='utf-8') as file:\n",
    "        return [json.loads(line) for line in file]\n",
    "\n",
    "def deduplicate_jsonl(json_objects):\n",
    "    \"\"\"De-duplicate a list of JSON objects based on their string representation.\"\"\"\n",
    "    unique_objects = set()\n",
    "    deduplicated_list = []\n",
    "    for obj in json_objects:\n",
    "        obj_str = json.dumps(obj, sort_keys=True)\n",
    "        if obj_str not in unique_objects:\n",
    "            unique_objects.add(obj_str)\n",
    "            deduplicated_list.append(obj)\n",
    "    return deduplicated_list\n",
    "\n",
    "def shuffle_and_write_jsonl(json_objects, output_path):\n",
    "    \"\"\"Shuffle a list of JSON objects and write them to a JSON Lines file.\"\"\"\n",
    "    random.shuffle(json_objects)\n",
    "    with open(output_path, 'w', encoding='utf-8') as file:\n",
    "        for obj in json_objects:\n",
    "            file.write(json.dumps(obj) + '\\n')\n",
    "\n",
    "# Paths to your JSON Lines files\n",
    "input_files = [\"datasets/arxiv_qa3.jsonl\", \"datasets/arxiv_metadata_qa3.jsonl\", \"datasets/arxiv_summary3.jsonl\"]\n",
    "output_path = \"datasets/arxiv_sharegpt2.jsonl\"\n",
    "\n",
    "# Read, combine, and de-duplicate\n",
    "combined_json_objects = []\n",
    "for file_path in input_files:\n",
    "    combined_json_objects.extend(read_jsonl(file_path))\n",
    "combined_json_objects = deduplicate_jsonl(combined_json_objects)\n",
    "\n",
    "# Shuffle and write to file\n",
    "shuffle_and_write_jsonl(combined_json_objects, output_path)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Optional: convert to LMFlow format\n",
    "https://optimalscale.github.io/LMFlow/examples/DATASETS.html\n",
    "\n",
    "Pretraining dataset should be JSON file like"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "{\n",
    "  \"type\": \"text_only\",\n",
    "  \"instances\": [\n",
    "    {  \"text\": \"SAMPLE_TEXT_1\" },\n",
    "    {  \"text\": \"SAMPLE_TEXT_2\" },\n",
    "    {  \"text\": \"SAMPLE_TEXT_3\" },\n",
    "  ]\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and QA tuning dataset should look like"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "{\n",
    "  \"type\": \"text2text\",\n",
    "  \"instances\": [\n",
    "    {\n",
    "        \"input\": \"SAMPLE_INPUT_1\",\n",
    "        \"output\": \"SAMPLE_OUTPUT_1\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"SAMPLE_INPUT_2\",\n",
    "        \"output\": \"SAMPLE_OUTPUT_2\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"SAMPLE_INPUT_3\",\n",
    "        \"output\": \"SAMPLE_OUTPUT_3\",\n",
    "    },\n",
    "  ]\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert pretraining data\n",
    "import json\n",
    "import glob\n",
    "import random\n",
    "pretraining_files = glob.glob(\"datasets/pretrain/*.jsonl\")\n",
    "out_file = \"datasets/lmflow/cosmo_pretrain.json\"\n",
    "out_dict = {\"type\": \"text_only\", \"instances\": []}\n",
    "for file in pretraining_files:\n",
    "    with open(file, \"r\") as f:\n",
    "        for line in f:\n",
    "            out_dict[\"instances\"].append(json.loads(line))\n",
    "random.shuffle(out_dict[\"instances\"])\n",
    "with open(out_file, \"w\") as f:\n",
    "    json.dump(out_dict, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert qa tuning data\n",
    "import json\n",
    "import glob\n",
    "import random\n",
    "qa_files = glob.glob(\"datasets/qa_tune/*.jsonl\")\n",
    "out_file = \"datasets/lmflow/cosmo_qa_tune.json\"\n",
    "out_dict = {\"type\": \"text2text\", \"instances\": []}\n",
    "for file in qa_files:\n",
    "    with open(file, \"r\") as f:\n",
    "        format = None\n",
    "        for iline,line in enumerate(f):\n",
    "            if iline==0:\n",
    "                if line[\"conversations\"]\n",
    "            out_dict[\"instances\"].append(json.loads(line))\n",
    "random.shuffle(out_dict[\"instances\"])\n",
    "with open(out_file, \"w\") as f:\n",
    "    json.dump(out_dict, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: Train the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we have two options. We can either keep control of the training loop. To do this uncomment and run the following code. The other option is to train on the JSONL files with the `axolotl` package. The advantage of this is that it comes with a lot of bells and whistles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# OPTION 5.1: keep control of the training loop\n",
    "# fine_tune.fine_tune(\n",
    "#     pretrained_model_file_path=\"zephyr-7b-beta\",\n",
    "#     training_data=cleaned_json_file_path,\n",
    "#     lr=5e-5,\n",
    "#     gradient_clip=1.0,\n",
    "#     num_epochs=1,\n",
    "#     out_dir=\"zephyr-7b-beta_cosmosage_v1\",\n",
    "# )\n",
    "\n",
    "# OPTION 5.2: train using axolotl and its config.yml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Axolotl \n",
    "\n",
    "You can run \n",
    "```accelerate launch -m axolotl.cli.train config.yml --prepare_ds_only --debug```\n",
    "to see examples of what data your model is being finetuned on. It is useful for knowing the exact prompt template to use during inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize loss during training\n",
    "import plot_tf_log\n",
    "v16 = plot_tf_log.most_recent_log(\"mistral_cosmosage_v16\")\n",
    "v14 = plot_tf_log.most_recent_log(\"mistral_cosmosage_v14\")\n",
    "v15 = plot_tf_log.most_recent_log(\"mistral_cosmosage_v15\")\n",
    "plot_tf_log.plot_loss([v15], plot_type=\"detailed\", detailed_pts_per_eval=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6: Evaluate the fine-tuned model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = \"models/cosmosage_v2/\"\n",
    "\n",
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "device = torch.device(\"cuda\")\n",
    "model = AutoModelForCausalLM.from_pretrained(model_path).to(device, dtype=torch.bfloat16)\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_path)\n",
    "\n",
    "def ask_cosmosage(question):\n",
    "    prompt = f\"You are cosmosage, an AI programmed to provide excellent and detailed answers to the user's question. You are an expert cosmology assistant, able to answer questions on the cosmic microwave background, galaxy formation, large scale structure, theoretical cosmology, inflation, big bang nucleosynthesis, cosmology instrumentation, and other related topics. Please assume the user is fluent in scientific terminology. Elaborate where possible to give a complete answer. If you do not know, say you do not know.▁ USER: {question}▁ ASSISTANT:\"\n",
    "    input_ids = tokenizer.encode(prompt, return_tensors=\"pt\").to(device)\n",
    "    generated_ids = model.generate(input_ids, max_length=1024, do_sample=True, temperature=0.7, top_k=None, pad_token_id=tokenizer.eos_token_id)\n",
    "    generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)\n",
    "    answer = generated_text.split(\"ASSISTANT: \")[-1]\n",
    "    return answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(ask_cosmosage(\"What is the temperature of the CMB according to Fixsen (2009)? What datasets did he use to derive the value?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(ask_cosmosage(\"What is Digital Active Nulling?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(ask_cosmosage(\"Explain the ISW effect.\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(ask_cosmosage(\"How does the time of matter-radiation equality affect the damping tail?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(ask_cosmosage(\"Explain how one would calculate the helium fraction at the surface of last scattering.\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 7: Push model to huggingface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import HfApi\n",
    "api = HfApi()\n",
    "api.upload_folder(\n",
    "    folder_path=\"models/cosmosage_v2\",\n",
    "    repo_id=\"tijmen2/cosmosage_v2\",\n",
    "    repo_type=\"model\",\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}