example_related_work_generation_full.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1d41424",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31034314",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "from time import sleep\n",
    "from copy import deepcopy\n",
    "import json\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "\n",
    "import os\n",
    "import sys\n",
    "from glob import glob\n",
    "from util import *\n",
    "from data_util import scientific_sent_tokenize\n",
    "from joint_tagger import CorwaTagger\n",
    "\n",
    "import torch\n",
    "from transformers import (\n",
    "    AutoConfig,\n",
    "    AutoModelForSequenceClassification,\n",
    "    AutoTokenizer,\n",
    ")\n",
    "from prompting_util import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "499372d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "prefix = \"example_pdfs/\"\n",
    "keyword = ['related work']\n",
    "target_file_name = \"file_name.json\"\n",
    "keyword_relation = \"in\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c3ba08f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def link_and_update_citations(this_paper):\n",
    "    discovered_citations = {}\n",
    "    for bib_key, bib in this_paper[\"pdf_parse\"][\"bib_entries\"].items():\n",
    "        if len(bib[\"authors\"]) > 0:\n",
    "            key = bib[\"title\"].lower() + \"@\" + bib[\"authors\"][0][\"last\"].replace(\" \",\"\").lower()\n",
    "            if key in title_author_lookup:\n",
    "                if title_author_lookup[key][\"year\"] == \"\":\n",
    "                    title_author_lookup[key][\"year\"] = bib[\"year\"]\n",
    "                discovered_citations[title_author_lookup[key][\"paper_id\"]] = bib_key\n",
    "    citation_mapping = this_paper.get(\"discovered_citations\",{})\n",
    "    citation_mapping.update(discovered_citations)\n",
    "    this_paper[\"discovered_citations\"] = citation_mapping\n",
    "    return this_paper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "227e867b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_cited_papers_text(chronological_paper_ids, use_faceted_summary=True, include_usage=True, include_relationship=True, use_CTS=False):\n",
    "    cited_papers_text = \"\"\n",
    "    for i, cited_paper_id in enumerate(chronological_paper_ids):\n",
    "        cited_metadata = global_metadata[cited_paper_id]\n",
    "        if cited_paper_id != target_paper_id:\n",
    "            cited_year = \"(\" + str(cited_metadata[\"year\"]) + \")\" if cited_metadata[\"year\"] else \"\"\n",
    "            cited_marker = cited_metadata[\"author\"] + \" et al. \" + cited_year\n",
    "            cited_title = cited_metadata[\"title\"] + \" by \" + cited_marker\n",
    "            \n",
    "            if use_faceted_summary:\n",
    "                cited_papers_text += str(i+1)+\". \" + cited_title\n",
    "                cited_papers_text += \"\\n\" + cited_facted_summaries[cited_paper_id] + \"\\n\\n\" \n",
    "            else:\n",
    "                cited_papers_text += str(i+1)+\". \" + cited_title\n",
    "                cited_papers_text += \"\\nAbstract:\" + cited_abstracts[cited_paper_id] + \"\\n\\n\" \n",
    "            \n",
    "            if include_usage and cited_paper_id in citation_intent_summary:\n",
    "                cited_papers_text += \"[Usage] \"+citation_intent_summary[cited_paper_id] + \"\\n\" \n",
    "                cited_papers_text += \"\\n\"\n",
    "\n",
    "            if include_relationship and cited_paper_id in relationships_by_cited_paper and len(relationships_by_cited_paper) > 0:\n",
    "                cited_papers_text += \"How other papers cite it: \\n\"\n",
    "                for relation in relationships_by_cited_paper[cited_paper_id]:\n",
    "                    cited_papers_text += relation[\"response\"] + \"\\n\" \n",
    "                cited_papers_text += \"\\n\"\n",
    "                \n",
    "            if use_CTS and len(cited_text_spans[cited_paper_id]) > 0:\n",
    "                cited_papers_text += \"Potentially useful sentences from this paper: \\n\"\n",
    "                for section, sentence in cited_text_spans[cited_paper_id]:\n",
    "                    cited_papers_text += \"[\"+section+\"] \"+sentence + \"\\n\" \n",
    "                cited_papers_text += \"\\n\"\n",
    "    cited_papers_text = cited_papers_text.replace(target_marker, \"our paper\")\n",
    "    return cited_papers_text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36e2ae02",
   "metadata": {},
   "outputs": [],
   "source": [
    "def retrieve_citation_spans(citing_paper_id, cited_paper_id, span_types = [\"Dominant\", \"Reference\"]):\n",
    "    spans = []\n",
    "    if cited_paper_id in cited_jsons[citing_paper_id][\"discovered_citations\"]:\n",
    "        cited_bib_key = cited_jsons[citing_paper_id][\"discovered_citations\"][cited_paper_id]\n",
    "        if cited_bib_key in span_citation_mappings_lookup[citing_paper_id]:\n",
    "            all_mentions = span_citation_mappings_lookup[citing_paper_id][cited_bib_key]\n",
    "\n",
    "            for span_type in span_types:\n",
    "                for mention, citation_mark in all_mentions[span_type].items():\n",
    "                    pid, cid = mention.split(\"_\")\n",
    "                    this_paragraph = span_citation_mappings_dict[citing_paper_id+\"_\"+pid]\n",
    "                    start = this_paragraph[\"span_citation_mapping\"][int(cid)][\"char_start\"]\n",
    "                    end = this_paragraph[\"span_citation_mapping\"][int(cid)][\"char_end\"]\n",
    "                    span_text = this_paragraph[\"paragraph\"][start:end].replace(\"[BOS]\",\"\")\n",
    "                    spans.append({\n",
    "                        \"paragraph_id\": pid,\n",
    "                        \"citation_id\": cid,\n",
    "                        \"type\": span_type,\n",
    "                        \"marker\": citation_mark,\n",
    "                        \"text\": span_text, \n",
    "                    })\n",
    "    return spans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93dbdb0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_relationship_prompt(retrieved_spans, citing_paper_id, cited_paper_id):\n",
    "    concatenated_spans = \"\\n\".join([str(i+1)+\". \"+span[\"text\"] for i, span in enumerate(retrieved_spans)])\n",
    "    concatenated_markers = \" or \".join(list(set([span[\"marker\"] for span in retrieved_spans])))\n",
    "\n",
    "    citing_metadata = global_metadata[citing_paper_id]\n",
    "    citing_year = \"(\" + str(citing_metadata[\"year\"]) + \")\" if citing_metadata[\"year\"] else \"\"\n",
    "    citing_marker = citing_metadata[\"author\"] + \" et al. \" + citing_year\n",
    "    citing_title = citing_metadata[\"title\"] + \" by \" + citing_marker\n",
    "\n",
    "    cited_metadata = global_metadata[cited_paper_id]\n",
    "    cited_year = \"(\" + str(cited_metadata[\"year\"]) + \")\" if cited_metadata[\"year\"] else \"\"\n",
    "    cited_marker = cited_metadata[\"author\"] + \" et al. \" + cited_year\n",
    "    cited_title = cited_metadata[\"title\"] + \" by \" + cited_marker\n",
    "\n",
    "    citing_paper = \"Faceted summary of the citing paper, \" + citing_title + \": \\n\" + cited_facted_summaries[citing_paper_id]\n",
    "    cited_paper = \"Faceted summary of the cited paper, \" + cited_title + \": \\n\" + cited_facted_summaries[cited_paper_id]\n",
    "    citation_span_texts = \"Citation contexts that \"+citing_marker+ \" cites \" + cited_marker+\" (which is cited as \"+concatenated_markers+\"): \\n\" + concatenated_spans\n",
    "    prompt_question = \"Very briefly explain the relationship between \" + cited_marker +\" and \" + citing_marker+\". TLDR: \"\n",
    "    relationship_prompt = \"\\n\\n\".join([citing_paper, cited_paper, citation_span_texts, prompt_question])\n",
    "    return relationship_prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "595ebbea",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "gpt_tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d244213",
   "metadata": {},
   "outputs": [],
   "source": [
    "argparser = ArgumentParser()\n",
    "argparser.add_argument('--repfile', type=str, default = \"allenai/scibert_scivocab_uncased\", help=\"Word embedding file\")\n",
    "argparser.add_argument('--dropout', type=float, default=0, help=\"embedding_dropout rate\")\n",
    "argparser.add_argument('--bert_dim', type=int, default=768, help=\"bert_dimension\")\n",
    "argparser.add_argument('--MAX_SENT_LEN', type=int, default=512)\n",
    "argparser.add_argument('--checkpoint', type=str, default = \"joint_tagger_train_scibert_final.model\")\n",
    "argparser.add_argument('--batch_size', type=int, default=1) # roberta-large: 2; bert: 8\n",
    "args = argparser.parse_args()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cea0fba",
   "metadata": {},
   "outputs": [],
   "source": [
    "target_paper_name = prefix+target_file_name\n",
    "with open(target_paper_name) as f:\n",
    "    target_json = json.load(f)\n",
    "# Exclude target paper's gold related work section.\n",
    "excluded_body_text = []\n",
    "for paragraph in target_json[\"pdf_parse\"][\"body_text\"]:\n",
    "    if not is_related_work_section(keyword, paragraph[\"section\"].lower(), keyword_relation):\n",
    "        excluded_body_text.append(paragraph)\n",
    "target_json_no_related_work = deepcopy(target_json)\n",
    "target_json_no_related_work[\"pdf_parse\"][\"body_text\"] = excluded_body_text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d1c0cef",
   "metadata": {},
   "outputs": [],
   "source": [
    "cited_jsons = {}\n",
    "for json_name in glob(prefix+\"*.json\"):\n",
    "    ID = \".\".join(json_name.split(\"/\")[-1].split(\".\")[:-1])\n",
    "    with open(json_name) as f:\n",
    "        paper = json.load(f)\n",
    "        try:\n",
    "            paper_id, paper = create_paper_id(paper, ID)\n",
    "            cited_jsons[paper_id] = paper\n",
    "        except:\n",
    "            pass\n",
    "target_paper_id, target_paper = create_paper_id(target_json_no_related_work)\n",
    "cited_jsons[target_paper_id] = target_paper\n",
    "target_paper_json = cited_jsons[target_paper_id]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85bd1e11",
   "metadata": {},
   "outputs": [],
   "source": [
    "pdf_parse_jsons = {paper_id: paper[\"pdf_parse\"]  for paper_id, paper in cited_jsons.items()}\n",
    "paragraphs = {}\n",
    "for paper_id, paper in pdf_parse_jsons.items():\n",
    "    for pi, para in enumerate(paper[\"body_text\"]):\n",
    "        paragraph_id = paper_id + \"_\" + str(pi)\n",
    "        paragraphs[paragraph_id] = \" \".join(scientific_sent_tokenize(para[\"text\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e20aa10f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "joint_tagger_tokenizer = AutoTokenizer.from_pretrained(\"allenai/scibert_scivocab_uncased\")\n",
    "additional_special_tokens = {'additional_special_tokens': ['[BOS]']}\n",
    "joint_tagger_tokenizer.add_special_tokens(additional_special_tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c2cf289",
   "metadata": {},
   "outputs": [],
   "source": [
    "tagger = CorwaTagger(joint_tagger_tokenizer, device, args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34e0bfd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_span_citation_mappings = tagger.run_prediction(paragraphs, pdf_parse_jsons)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb6c6107",
   "metadata": {},
   "outputs": [],
   "source": [
    "span_citation_mappings_dict = {paragraph[\"id\"]: paragraph for paragraph in all_span_citation_mappings} "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a0c51af",
   "metadata": {},
   "outputs": [],
   "source": [
    "span_citation_mappings_lookup = {}\n",
    "for paragraph in all_span_citation_mappings:\n",
    "    paper_id, pid = paragraph[\"id\"].split(\"_\")\n",
    "    this_paper = span_citation_mappings_lookup.get(paper_id,{})\n",
    "    for si, span in enumerate(paragraph[\"span_citation_mapping\"]):\n",
    "        for span_type, citations in span[\"span_citation_mapping\"].items():\n",
    "            for citation_mark, bib_key in citations.items():\n",
    "                this_cited_paper = this_paper.get(bib_key,{\"Dominant\":{}, \"Reference\":{}})\n",
    "                this_cited_paper[span_type][pid+\"_\"+str(si)] = citation_mark\n",
    "                this_paper[bib_key] = this_cited_paper\n",
    "    span_citation_mappings_lookup[paper_id] = this_paper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d58da9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "global_metadata = {}\n",
    "for paper_id, paper in cited_jsons.items():\n",
    "    if len(paper[\"authors\"]) > 0:\n",
    "        global_metadata[paper_id] = {\n",
    "            \"paper_id\": paper_id,\n",
    "            \"title\": paper[\"title\"],\n",
    "            \"author\": paper[\"authors\"][0][\"last\"].replace(\" \",\"\"),\n",
    "            \"year\": paper[\"year\"],\n",
    "        }\n",
    "    else:\n",
    "        global_metadata[paper_id] = {\n",
    "            \"paper_id\": paper_id,\n",
    "            \"title\": paper[\"title\"],\n",
    "            \"author\": \"Unknown\",\n",
    "            \"year\": paper[\"year\"],\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b04459db",
   "metadata": {},
   "outputs": [],
   "source": [
    "title_author_lookup = {}\n",
    "for paper_id, paper in global_metadata.items():\n",
    "    key = paper[\"title\"].lower() + \"@\" + paper[\"author\"].lower().replace(\" \",\"\")\n",
    "    title_author_lookup[key] = paper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2c3605a",
   "metadata": {},
   "outputs": [],
   "source": [
    "for key, paper in cited_jsons.items():\n",
    "    this_paper = link_and_update_citations(paper)\n",
    "    if global_metadata[key][\"year\"] == \"\":\n",
    "        global_metadata[key][\"year\"] = key.split(\"@\")[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "023bc7ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(global_metadata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a084a90",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "api_key = \"\"\n",
    "openai.organization = \"\"\n",
    "openai.api_key = api_key\n",
    "openai.Model.list()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "756ec17e",
   "metadata": {},
   "outputs": [],
   "source": [
    "cited_facted_summaries = {}\n",
    "while len(cited_jsons) > len(cited_facted_summaries):\n",
    "    for key in tqdm(cited_jsons.keys()):\n",
    "        if key in cited_facted_summaries:\n",
    "            continue\n",
    "        cited_json = cited_jsons[key]\n",
    "        TAIC = \"\"\n",
    "        TAIC += \"[Title] \" + cited_json[\"title\"] + \" \\n\"\n",
    "        if cited_json[\"pdf_parse\"][\"abstract\"]:\n",
    "            for paragraph in cited_json[\"pdf_parse\"][\"abstract\"]:\n",
    "                TAIC += \"[Abstract] \" + paragraph[\"text\"] + \" \\n\"\n",
    "        for paragraph in cited_json[\"pdf_parse\"][\"body_text\"]:\n",
    "            if \"intro\" in paragraph[\"section\"].lower() or \"conclusion\" in paragraph[\"section\"].lower():\n",
    "                if not has_tokenization_error(paragraph[\"text\"]):\n",
    "                    TAIC += \"[\"+paragraph[\"section\"]+\"] \" + paragraph[\"text\"] + \" \\n\"\n",
    "\n",
    "        facet_prompt = TAIC + \"What are the objective, method, findings, contributions and keywords of the paper above? Answer in the format of \\n Objective: XXX. \\n Method: XXX. \\n Findings: XXX. \\n Contribution: XXX. \\n Keywords: A; B; C.\"\n",
    "        try:\n",
    "            response = openai.ChatCompletion.create(\n",
    "              model=\"gpt-3.5-turbo-0301\",\n",
    "              messages=[\n",
    "                    {\"role\": \"user\", \"content\": facet_prompt}\n",
    "                ],\n",
    "              temperature = 0,\n",
    "            )\n",
    "            cited_facted_summaries[key] = response[\"choices\"][0][\"message\"][\"content\"]\n",
    "        except:\n",
    "            print(\"Failed\", key)\n",
    "        #sleep(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c287c1fa",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "relationships = {}\n",
    "new_entry = -1\n",
    "while new_entry != 0:\n",
    "    new_entry = 0\n",
    "    for cited_paper_id in tqdm(cited_jsons.keys()):\n",
    "        for citing_paper_id in cited_jsons.keys():\n",
    "            if cited_paper_id + \"+\" + citing_paper_id in relationships:\n",
    "                continue\n",
    "            retrieved_spans = retrieve_citation_spans(citing_paper_id, cited_paper_id)\n",
    "            if len(retrieved_spans) > 0:\n",
    "                relationship_prompt = make_relationship_prompt(retrieved_spans, citing_paper_id, cited_paper_id)\n",
    "                try:\n",
    "                    response = openai.ChatCompletion.create(\n",
    "                      model=\"gpt-3.5-turbo-0301\",\n",
    "                      messages=[\n",
    "                            {\"role\": \"user\", \"content\": relationship_prompt}\n",
    "                        ],\n",
    "                      temperature = 0,\n",
    "                    )\n",
    "                    response_text = response[\"choices\"][0][\"message\"][\"content\"]\n",
    "                    this_pair = {\n",
    "                        \"citing_paper\": citing_paper_id,\n",
    "                        \"cited_paper\": cited_paper_id,\n",
    "                        \"prompt\": relationship_prompt, \n",
    "                        \"response\": response_text\n",
    "                    }\n",
    "                    relationships[cited_paper_id + \"+\" + citing_paper_id] = this_pair\n",
    "                    #print(this_pair[\"response\"])\n",
    "                    new_entry += 1\n",
    "                except:\n",
    "                    print(\"#\"*30)\n",
    "                    print(citing_paper_id)\n",
    "                    print(cited_paper_id)\n",
    "                    print(\"Failed!\")\n",
    "                #sleep(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c66835d",
   "metadata": {},
   "outputs": [],
   "source": [
    "relationships_by_cited_paper = {}\n",
    "for key, relation in relationships.items():\n",
    "    this_paper = relationships_by_cited_paper.get(relation[\"cited_paper\"],[])\n",
    "    this_paper.append(relation)\n",
    "    relationships_by_cited_paper[relation[\"cited_paper\"]] = this_paper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66adcc32",
   "metadata": {},
   "outputs": [],
   "source": [
    "citation_spans_by_cited_paper = {}\n",
    "for citing_paper_id in cited_jsons.keys():\n",
    "    for cited_paper_id in cited_jsons.keys():\n",
    "        retrieved_spans = retrieve_citation_spans(citing_paper_id, cited_paper_id)\n",
    "        if len(retrieved_spans) > 0:\n",
    "            this_paper = citation_spans_by_cited_paper.get(cited_paper_id,{})\n",
    "            this_paper[citing_paper_id] = retrieved_spans\n",
    "            citation_spans_by_cited_paper[cited_paper_id] = this_paper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3374add",
   "metadata": {},
   "outputs": [],
   "source": [
    "citation_intent_summary = {}\n",
    "while len(citation_intent_summary) < len(relationships_by_cited_paper):\n",
    "    for cited_paper_id, relations in tqdm(relationships_by_cited_paper.items()):\n",
    "        if cited_paper_id in citation_intent_summary:\n",
    "            continue\n",
    "        cited_metadata = global_metadata[cited_paper_id]\n",
    "        cited_year = \"(\" + str(cited_metadata[\"year\"]) + \")\" if cited_metadata[\"year\"] else \"\"\n",
    "        cited_marker = cited_metadata[\"author\"] + \" et al. \" + cited_year\n",
    "        cited_title = cited_metadata[\"title\"] + \" by \" + cited_marker\n",
    "\n",
    "        summary = cited_facted_summaries[cited_paper_id]\n",
    "\n",
    "        concatenated_relation = []\n",
    "        for relation in relations:\n",
    "            relation_text = relation[\"response\"]\n",
    "            citing_paper_id = relation[\"citing_paper\"]\n",
    "            if citing_paper_id in citation_spans_by_cited_paper[cited_paper_id]:\n",
    "                examples = []\n",
    "                for ci, citation in enumerate(citation_spans_by_cited_paper[cited_paper_id][citing_paper_id]):\n",
    "                    examples.append(str(ci+1)+\". \" + citation[\"text\"])\n",
    "                concatenated_relation.append(relation_text+\"\\nExample citation fragments: \\n\" + \"\\n\".join(examples))\n",
    "        concatenated_relation = \"\\n\\n\".join(concatenated_relation)\n",
    "\n",
    "        #citation_intent_prompt = \"Faceted summary of \" + cited_title +\":\\n\"\n",
    "        #citation_intent_prompt += summary + \"\\n\\n\"\n",
    "        #citation_intent_prompt += \"How other papers cite it: \\n\" + concatenated_relation + \"\\n\\n\"\n",
    "        # Faceted summary is a distractor!!\n",
    "        citation_intent_prompt = \"How other papers cite \"+cited_marker+\": \\n\" + concatenated_relation + \"\\n\\n\"\n",
    "        citation_intent_prompt += \"Very briefly answer what \"+cited_marker+\\\n",
    "        \" is mostly known for, and the common citation intent. \"\n",
    "        citation_intent_prompt += \"Hint: pay attention to how \"+cited_marker+\" is referred by the citing papers. \"\n",
    "        citation_intent_prompt += 'Answer in the format of \"'+cited_marker+' is known for XXX and it is cited for YYY\". TLDR: '\n",
    "\n",
    "        try:\n",
    "            response = openai.ChatCompletion.create(\n",
    "                      model=\"gpt-3.5-turbo-0301\",\n",
    "                      messages=[\n",
    "                            {\"role\": \"user\", \"content\": citation_intent_prompt}\n",
    "                        ],\n",
    "                      temperature = 0,\n",
    "                    )\n",
    "            response_text = response[\"choices\"][0][\"message\"][\"content\"]\n",
    "            #print(response_text)\n",
    "            citation_intent_summary[cited_paper_id] = response_text\n",
    "        except:\n",
    "            print(cited_paper_id, \"failed!\")\n",
    "        #sleep(10)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e822eaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "gold_related_work = \"\"\n",
    "for paragraph in target_json[\"pdf_parse\"][\"body_text\"]:\n",
    "    if is_related_work_section(keyword, paragraph[\"section\"].lower(), keyword_relation):\n",
    "        gold_related_work += paragraph[\"text\"] + \" \\n\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6648d064",
   "metadata": {},
   "outputs": [],
   "source": [
    "faceted_summary = cited_facted_summaries[target_paper_id]\n",
    "title = global_metadata[target_paper_id][\"title\"]\n",
    "high_level_idea_prompt = \"Our title: \" + title + \"\\nFaceted summary of our paper: \" + faceted_summary + \"\\n\\nWrite a short summary of the main idea of the following related work section paragraphs. Ignore citations.\\n\\n\" + gold_related_work"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "425e880e",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = openai.ChatCompletion.create(\n",
    "  model=\"gpt-3.5-turbo-0301\",\n",
    "  messages=[\n",
    "        {\"role\": \"user\", \"content\": high_level_idea_prompt}\n",
    "    ],\n",
    "  temperature = 0,\n",
    ")\n",
    "main_idea = response[\"choices\"][0][\"message\"][\"content\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54ee6f8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(main_idea)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90c6d255",
   "metadata": {},
   "outputs": [],
   "source": [
    "TAIC = \"\"\n",
    "TAIC += \"[Title] \" + target_paper_json[\"title\"] + \" \\n\"\n",
    "if target_paper_json[\"pdf_parse\"][\"abstract\"]:\n",
    "    for paragraph in target_paper_json[\"pdf_parse\"][\"abstract\"]:\n",
    "        TAIC += \"[Abstract] \" + paragraph[\"text\"] + \" \\n\"\n",
    "for paragraph in target_paper_json[\"pdf_parse\"][\"body_text\"]:\n",
    "    if \"intro\" in paragraph[\"section\"].lower() or \"conclusion\" in paragraph[\"section\"].lower():\n",
    "        TAIC += \"[\"+paragraph[\"section\"]+\"] \" + paragraph[\"text\"] + \" \\n\"\n",
    "target_paper_TAIC = TAIC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23fe7a63",
   "metadata": {},
   "outputs": [],
   "source": [
    "chronological_paper_ids = sorted(list(global_metadata.keys()), key=lambda x: x.split(\"@\")[1], reverse=False)\n",
    "target_paper = global_metadata[target_paper_id]\n",
    "target_marker = target_paper[\"author\"]+\" et al. (\" + str(target_paper[\"year\"]) + \")\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "798dd5fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "cited_papers_text = prepare_cited_papers_text(chronological_paper_ids, include_usage=True, include_relationship=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60b6d2ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "intro_instruction = \"We have finished writing the title, abstract, introduction and conclusion section of our NLP paper as follows: \\n\"\n",
    "semi_automatic_generation_instruction = \"\"\"\n",
    "However, the related work section is still missing.\n",
    "Write our related work section that concisely cites all the following papers in a natural way using all of the main ideas as the main story.\n",
    "Keep it short, e.g. a few paragraphs at most. Make sure the related work section do not conflict with the sections already written.\n",
    "You can freely reorder the cited papers to adapt to the main ideas.\n",
    "Pay extra attention to [Usage] which indicates how each work is cited by other work. \\n\n",
    "\"\"\"\n",
    "semi_automatic_generation_prompt = intro_instruction + target_paper_TAIC + semi_automatic_generation_instruction +\\\n",
    "\"\\n Main idea of our related work section: \" + main_idea + \"\\n\\nList of cited papers: \"+ cited_papers_text + \"\\nOur related work section: \" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5735b97a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(semi_automatic_generation_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fe11d80",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(gpt_tokenizer(semi_automatic_generation_prompt).input_ids) <= 8000\n",
    "response = openai.ChatCompletion.create(\n",
    "  model=\"gpt-4-0314\",\n",
    "  messages=[\n",
    "        #{\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": semi_automatic_generation_prompt}\n",
    "    ],\n",
    "  temperature = 0,\n",
    ")\n",
    "semi_automatic_related_work_gpt4 = response[\"choices\"][0][\"message\"][\"content\"]\n",
    "print(semi_automatic_related_work_gpt4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c5374c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "intro_instruction = \"We have finished writing the title, abstract, introduction and conclusion section of our NLP paper as follows: \\n\"\n",
    "fully_automatic_generation_instruction = \"\"\"\n",
    "However, the related work section is still missing.\n",
    "Write our related work section that concisely cites all the following papers in a natural way.\n",
    "Keep it short, e.g. a few paragraphs at most. Stick the topic to the main topic of our paper.\n",
    "Pay extra attention to [Usage] which indicates how each work is cited by other work. \\n\n",
    "\"\"\"\n",
    "fully_automatic_generation_prompt = intro_instruction + target_paper_TAIC + fully_automatic_generation_instruction + \"List of cited papers: \"+ cited_papers_text + \"\\nOur related work section: \"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21fbc72d",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(gpt_tokenizer(fully_automatic_generation_prompt).input_ids) <= 8000\n",
    "\n",
    "response = openai.ChatCompletion.create(\n",
    "  model=\"gpt-4-0314\",\n",
    "  messages=[\n",
    "        #{\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": fully_automatic_generation_prompt}\n",
    "    ],\n",
    "  temperature = 0,\n",
    ")\n",
    "fully_automatic_related_work_gpt4 = response[\"choices\"][0][\"message\"][\"content\"]\n",
    "print(fully_automatic_related_work_gpt4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8c8b480",
   "metadata": {},
   "outputs": [],
   "source": [
    "cited_papers_text_no_usage = prepare_cited_papers_text(chronological_paper_ids, include_usage=False, include_relationship=True)\n",
    "semi_automatic_generation_prompt_no_usage = intro_instruction + target_paper_TAIC + semi_automatic_generation_instruction +\\\n",
    "\"\\n Main idea of our related work section: \" + main_idea + \"\\n\\nList of cited papers: \"+ cited_papers_text_no_usage + \"\\nOur related work section: \" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f37d2e7",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "response = openai.ChatCompletion.create(\n",
    "  model=\"gpt-4-0314\",\n",
    "  messages=[\n",
    "        #{\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": semi_automatic_generation_prompt_no_usage}\n",
    "    ],\n",
    "  temperature = 0,\n",
    ")\n",
    "semi_automatic_related_work_gpt4_no_usage = response[\"choices\"][0][\"message\"][\"content\"]\n",
    "print(semi_automatic_related_work_gpt4_no_usage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eea1108f",
   "metadata": {},
   "outputs": [],
   "source": [
    "cited_papers_text_no_relationship = prepare_cited_papers_text(chronological_paper_ids, include_usage=True, include_relationship=False)\n",
    "semi_automatic_generation_prompt_no_relationship = intro_instruction + target_paper_TAIC + semi_automatic_generation_instruction +\\\n",
    "\"\\n Main idea of our related work section: \" + main_idea + \"\\n\\nList of cited papers: \"+ cited_papers_text_no_relationship + \"\\nOur related work section: \" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "394999c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = openai.ChatCompletion.create(\n",
    "  model=\"gpt-4-0314\",\n",
    "  messages=[\n",
    "        #{\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": semi_automatic_generation_prompt_no_relationship}\n",
    "    ],\n",
    "  temperature = 0,\n",
    ")\n",
    "semi_automatic_related_work_gpt4_no_relationship = response[\"choices\"][0][\"message\"][\"content\"]\n",
    "print(semi_automatic_related_work_gpt4_no_relationship)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b70e66e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "intro_instruction = \"We have finished writing the title, abstract, introduction and conclusion section of our NLP paper. \"\n",
    "semi_automatic_generation_instruction = \"\"\"\n",
    "However, the related work section is still missing.\n",
    "Write our related work section that concisely cites all the following papers in a natural way using all of the main ideas as the main story.\n",
    "Keep it short, e.g. a few paragraphs at most. Make sure the related work section do not conflict with the sections already written.\n",
    "You can freely reorder the cited papers to adapt to the main ideas.\n",
    "Pay extra attention to [Usage] which indicates how each work is cited by other work. \\n\n",
    "\"\"\"\n",
    "semi_automatic_generation_prompt_no_TAIC = intro_instruction + semi_automatic_generation_instruction +\\\n",
    "\"\\n Main idea of our related work section: \" + main_idea + \"\\n\\nList of cited papers: \"+ cited_papers_text + \"\\nOur related work section: \" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4cd330b",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = openai.ChatCompletion.create(\n",
    "  model=\"gpt-4-0314\",\n",
    "  messages=[\n",
    "        #{\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": semi_automatic_generation_prompt_no_TAIC}\n",
    "    ],\n",
    "  temperature = 0,\n",
    ")\n",
    "semi_automatic_related_work_gpt4_no_TAIC = response[\"choices\"][0][\"message\"][\"content\"]\n",
    "print(semi_automatic_related_work_gpt4_no_TAIC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a84938f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_paragraphs = {}\n",
    "pi = 0\n",
    "for paragraph in semi_automatic_related_work_gpt4.split(\"\\n\"):\n",
    "    if paragraph:\n",
    "        output_paragraphs[\"predicted_\"+str(pi)] = \" \".join(scientific_sent_tokenize(paragraph))\n",
    "        pi+=1\n",
    "        \n",
    "pseudo_related_work_json = {\n",
    "    \"paper_id\": \"predicted\",\n",
    "    \"bib_entries\": {},\n",
    "    \"body_text\": []\n",
    "}\n",
    "for para_id, paragraph in output_paragraphs.items():\n",
    "    pseudo_related_work_json[\"body_text\"].append({\n",
    "        \"section\": para_id,\n",
    "        \"text\": paragraph,\n",
    "        \"cite_spans\": [],\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59615815",
   "metadata": {},
   "outputs": [],
   "source": [
    "generated_span_citation_mappings = tagger.run_prediction(output_paragraphs, {\"predicted\": pseudo_related_work_json})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c508b51",
   "metadata": {},
   "outputs": [],
   "source": [
    "for paragraph in generated_span_citation_mappings:\n",
    "    for span in paragraph[\"span_citation_mapping\"]:\n",
    "        for citation_type, citation in span[\"span_citation_mapping\"].items():\n",
    "            for marker in citation.keys():\n",
    "                paper_id = match_paper_id(marker, global_metadata)\n",
    "                if paper_id:\n",
    "                    citation[marker] = paper_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cecbbce3",
   "metadata": {},
   "outputs": [],
   "source": [
    "cited_paper_sentences = {}\n",
    "for paper_id, cited_json in cited_jsons.items():\n",
    "    sentences = []\n",
    "    for paragraph in cited_json[\"pdf_parse\"][\"body_text\"]:\n",
    "        these_sentences = scientific_sent_tokenize(paragraph[\"text\"],add_bos_token=False)\n",
    "        for sentence in these_sentences:\n",
    "            sentences.append((paragraph[\"section\"], sentence))\n",
    "    cited_paper_sentences[paper_id] = sentences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cb843a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "intro_instruction = \"We have finished writing the title, abstract, introduction and conclusion section of our NLP paper as follows: \\n\"\n",
    "semi_automatic_generation_instruction = \"\"\"\n",
    "However, the related work section is still missing.\n",
    "Write our related work section that concisely cites the following papers in a natural way using all of the main ideas as the main story.\n",
    "Keep it short, e.g. 3 paragraphs at most. Make sure the related work section do not conflict with the sections already written.\n",
    "You can freely reorder the cited papers to adapt to the main ideas.\n",
    "Pay extra attention to [Usage] which indicates how each work is cited by other work. \\n\n",
    "\"\"\"\n",
    " \n",
    "\n",
    "for k in [10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0]:\n",
    "    cited_text_spans = {paper_id: [] for paper_id in cited_paper_sentences.keys()}\n",
    "    for paragraph in generated_span_citation_mappings:\n",
    "        text = paragraph[\"paragraph\"]\n",
    "        for span in paragraph[\"span_citation_mapping\"]:\n",
    "            start = span[\"char_start\"]\n",
    "            end = span[\"char_end\"]\n",
    "            span_text = text[start:end]\n",
    "            for citation_type, citation in span[\"span_citation_mapping\"].items():\n",
    "                for marker, paper_id in citation.items():\n",
    "                    if paper_id is not None and cited_paper_sentences[paper_id] is not None:\n",
    "                        retrieved_sentences = rouge_retrieval(span_text, cited_paper_sentences[paper_id],k=k)\n",
    "                        cited_text_spans[paper_id].extend(retrieved_sentences)\n",
    "\n",
    "    cited_papers_text_CTS = prepare_cited_papers_text(chronological_paper_ids, use_faceted_summary=True, include_usage=True, include_relationship=True, use_CTS=True)\n",
    "    semi_automatic_generation_prompt_CTS = intro_instruction + target_paper_TAIC + semi_automatic_generation_instruction +\\\n",
    "    \"\\n Main idea of our related work section: \" + main_idea + \"\\n\\n List of cited papers: \"+\\\n",
    "    cited_papers_text_CTS + \"\\nOur related work section that cites ALL of the listed papers above: \" \n",
    "    if len(gpt_tokenizer(semi_automatic_generation_prompt_CTS).input_ids) <= 8000:\n",
    "        print(\"k=\"+str(k))\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c30528a4",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(semi_automatic_generation_prompt_CTS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "956f902b",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(gpt_tokenizer(semi_automatic_generation_prompt_CTS).input_ids) <= 8000\n",
    "\n",
    "response = openai.ChatCompletion.create(\n",
    "  model=\"gpt-4-0314\",\n",
    "  messages=[\n",
    "        #{\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": semi_automatic_generation_prompt_CTS}\n",
    "    ],\n",
    "  temperature = 0,\n",
    ")\n",
    "related_work_text_CTS = response[\"choices\"][0][\"message\"][\"content\"]\n",
    "print(related_work_text_CTS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "981a8854",
   "metadata": {},
   "outputs": [],
   "source": [
    "cited_abstracts = {}\n",
    "for key in cited_jsons.keys():\n",
    "    if key in cited_abstracts:\n",
    "        continue\n",
    "    cited_json = cited_jsons[key]\n",
    "    if cited_json[\"pdf_parse\"][\"abstract\"]:\n",
    "        cited_abstracts[key] = cited_json[\"pdf_parse\"][\"abstract\"][0][\"text\"]\n",
    "    else:\n",
    "        TAIC = \"\"\n",
    "        for paragraph in cited_json[\"pdf_parse\"][\"body_text\"]:\n",
    "            if \"intro\" in paragraph[\"section\"].lower():\n",
    "                TAIC += paragraph[\"text\"] + \" \\n\"\n",
    "        cited_abstracts[key] = TAIC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8db149e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "intro_instruction = \"We have finished writing the title, abstract, introduction and conclusion section of our NLP paper as follows: \\n\"\n",
    "semi_automatic_generation_instruction = \"\"\"\n",
    "However, the related work section is still missing.\n",
    "Write our related work section that concisely cites all the following papers in a natural way using all of the main ideas as the main story.\n",
    "Keep it short, e.g. a few paragraphs at most. Make sure the related work section do not conflict with the sections already written.\n",
    "You can freely reorder the cited papers to adapt to the main ideas.\n",
    "Pay extra attention to [Usage] which indicates how each work is cited by other work. \\n\n",
    "\"\"\"\n",
    "cited_papers_text_abstract = prepare_cited_papers_text(chronological_paper_ids, use_faceted_summary = False, include_usage=True, include_relationship=True)\n",
    "semi_automatic_generation_prompt_abstract = intro_instruction + target_paper_TAIC + semi_automatic_generation_instruction +\\\n",
    "\"\\n Main idea of our related work section: \" + main_idea + \"\\n\\nList of cited papers: \"+ cited_papers_text_abstract + \"\\nOur related work section: \" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7714788e",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(gpt_tokenizer(semi_automatic_generation_prompt_abstract).input_ids) <= 8000\n",
    "\n",
    "response = openai.ChatCompletion.create(\n",
    "  model=\"gpt-4-0314\",\n",
    "  messages=[\n",
    "        #{\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": semi_automatic_generation_prompt_abstract}\n",
    "    ],\n",
    "  temperature = 0,\n",
    ")\n",
    "semi_automatic_related_work_gpt4_abstract = response[\"choices\"][0][\"message\"][\"content\"]\n",
    "print(semi_automatic_related_work_gpt4_abstract)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29f9771c",
   "metadata": {},
   "outputs": [],
   "source": [
    "intro_instruction = \"We have finished writing the title, abstract, introduction and conclusion section of our NLP paper as follows: \\n\"\n",
    "semi_automatic_generation_instruction_no_usage_relationship = \"\"\"\n",
    "However, the related work section is still missing.\n",
    "Write our related work section that concisely cites all the following papers in a natural way using all of the main ideas as the main story.\n",
    "Keep it short, e.g. a few paragraphs at most. Make sure the related work section do not conflict with the sections already written.\n",
    "You can freely reorder the cited papers to adapt to the main ideas.\\n\n",
    "\"\"\"\n",
    "cited_papers_text_no_usage_relationship = prepare_cited_papers_text(chronological_paper_ids, include_usage=False, include_relationship=False)\n",
    "semi_automatic_generation_prompt_no_usage_relationship = intro_instruction + target_paper_TAIC + semi_automatic_generation_instruction_no_usage_relationship +\\\n",
    "\"\\n Main idea of our related work section: \" + main_idea + \"\\n\\nList of cited papers: \"+ cited_papers_text_no_usage_relationship + \"\\nOur related work section: \" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2091c2e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = openai.ChatCompletion.create(\n",
    "  model=\"gpt-4-0314\",\n",
    "  messages=[\n",
    "        #{\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": semi_automatic_generation_prompt_no_usage_relationship}\n",
    "    ],\n",
    "  temperature = 0,\n",
    ")\n",
    "semi_automatic_related_work_gpt4_no_usage_relationship = response[\"choices\"][0][\"message\"][\"content\"]\n",
    "print(semi_automatic_related_work_gpt4_no_usage_relationship)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba9169ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_paragraphs = {}\n",
    "pi = 0\n",
    "for paragraph in semi_automatic_related_work_gpt4_abstract.split(\"\\n\"):\n",
    "    if paragraph:\n",
    "        output_paragraphs[\"predicted_\"+str(pi)] = \" \".join(scientific_sent_tokenize(paragraph))\n",
    "        pi+=1\n",
    "        \n",
    "pseudo_related_work_json = {\n",
    "    \"paper_id\": \"predicted\",\n",
    "    \"bib_entries\": {},\n",
    "    \"body_text\": []\n",
    "}\n",
    "for para_id, paragraph in output_paragraphs.items():\n",
    "    pseudo_related_work_json[\"body_text\"].append({\n",
    "        \"section\": para_id,\n",
    "        \"text\": paragraph,\n",
    "        \"cite_spans\": [],\n",
    "    })\n",
    "    \n",
    "generated_span_citation_mappings = tagger.run_prediction(output_paragraphs, {\"predicted\": pseudo_related_work_json})\n",
    "\n",
    "for paragraph in generated_span_citation_mappings:\n",
    "    for span in paragraph[\"span_citation_mapping\"]:\n",
    "        for citation_type, citation in span[\"span_citation_mapping\"].items():\n",
    "            for marker in citation.keys():\n",
    "                paper_id = match_paper_id(marker, global_metadata)\n",
    "                if paper_id:\n",
    "                    citation[marker] = paper_id\n",
    "                    \n",
    "cited_paper_sentences = {}\n",
    "for paper_id, cited_json in cited_jsons.items():\n",
    "    sentences = []\n",
    "    for paragraph in cited_json[\"pdf_parse\"][\"body_text\"]:\n",
    "        these_sentences = scientific_sent_tokenize(paragraph[\"text\"],add_bos_token=False)\n",
    "        for sentence in these_sentences:\n",
    "            sentences.append((paragraph[\"section\"], sentence))\n",
    "    cited_paper_sentences[paper_id] = sentences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "faeb262e",
   "metadata": {},
   "outputs": [],
   "source": [
    "intro_instruction = \"We have finished writing the title, abstract, introduction and conclusion section of our NLP paper as follows: \\n\"\n",
    "semi_automatic_generation_instruction = \"\"\"\n",
    "However, the related work section is still missing.\n",
    "Write our related work section that concisely cites the following papers in a natural way using all of the main ideas as the main story.\n",
    "Keep it short, e.g. 3 paragraphs at most. Make sure the related work section do not conflict with the sections already written.\n",
    "You can freely reorder the cited papers to adapt to the main ideas.\n",
    "Pay extra attention to [Usage] which indicates how each work is cited by other work. \\n\n",
    "\"\"\"\n",
    " \n",
    "\n",
    "for k in [10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0]:\n",
    "    cited_text_spans = {paper_id: [] for paper_id in cited_paper_sentences.keys()}\n",
    "    for paragraph in generated_span_citation_mappings:\n",
    "        text = paragraph[\"paragraph\"]\n",
    "        for span in paragraph[\"span_citation_mapping\"]:\n",
    "            start = span[\"char_start\"]\n",
    "            end = span[\"char_end\"]\n",
    "            span_text = text[start:end]\n",
    "            for citation_type, citation in span[\"span_citation_mapping\"].items():\n",
    "                for marker, paper_id in citation.items():\n",
    "                    if paper_id is not None and cited_paper_sentences[paper_id] is not None:\n",
    "                        retrieved_sentences = rouge_retrieval(span_text, cited_paper_sentences[paper_id],k=k)\n",
    "                        cited_text_spans[paper_id].extend(retrieved_sentences)\n",
    "\n",
    "    cited_papers_text_abstract_CTS = prepare_cited_papers_text(chronological_paper_ids, use_faceted_summary=False, include_usage=True, include_relationship=True, use_CTS=True)\n",
    "    semi_automatic_generation_prompt_abstract_CTS = intro_instruction + target_paper_TAIC + semi_automatic_generation_instruction +\\\n",
    "    \"\\n Main idea of our related work section: \" + main_idea + \"\\n\\n List of cited papers: \"+\\\n",
    "    cited_papers_text_abstract_CTS + \"\\nOur related work section that cites ALL of the listed papers above: \" \n",
    "    if len(gpt_tokenizer(semi_automatic_generation_prompt_abstract_CTS).input_ids) <= 8000:\n",
    "        print(\"k=\"+str(k))\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b6dd146",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(gpt_tokenizer(semi_automatic_generation_prompt_abstract_CTS).input_ids) <= 8000\n",
    "\n",
    "response = openai.ChatCompletion.create(\n",
    "  model=\"gpt-4-0314\",\n",
    "  messages=[\n",
    "        #{\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": semi_automatic_generation_prompt_abstract_CTS}\n",
    "    ],\n",
    "  temperature = 0,\n",
    ")\n",
    "related_work_text_abstract_CTS = response[\"choices\"][0][\"message\"][\"content\"]\n",
    "print(related_work_text_abstract_CTS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02ae6f0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "intro_instruction = \"We have finished writing the title, abstract, introduction and conclusion section of our NLP paper as follows: \\n\"\n",
    "semi_automatic_generation_instruction = \"\"\"\n",
    "However, the related work section is still missing.\n",
    "Write our related work section that concisely cites all the following papers in a natural way using all of the main ideas as the main story.\n",
    "Keep it short, e.g. a few paragraphs at most. Make sure the related work section do not conflict with the sections already written.\n",
    "You can freely reorder the cited papers to adapt to the main ideas.\\n\n",
    "\"\"\"\n",
    "cited_papers_text_abstract_no_usage_relationship = prepare_cited_papers_text(chronological_paper_ids, use_faceted_summary = False, include_usage=False, include_relationship=False)\n",
    "semi_automatic_generation_prompt_abstract_no_usage_relationship = intro_instruction + target_paper_TAIC + semi_automatic_generation_instruction +\\\n",
    "\"\\n Main idea of our related work section: \" + main_idea + \"\\n\\nList of cited papers: \"+ cited_papers_text_abstract_no_usage_relationship + \"\\nOur related work section: \" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d70c26e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(gpt_tokenizer(semi_automatic_generation_prompt_abstract_no_usage_relationship).input_ids) <= 8000\n",
    "\n",
    "response = openai.ChatCompletion.create(\n",
    "  model=\"gpt-4-0314\",\n",
    "  messages=[\n",
    "        #{\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": semi_automatic_generation_prompt_abstract_no_usage_relationship}\n",
    "    ],\n",
    "  temperature = 0,\n",
    ")\n",
    "semi_automatic_related_work_gpt4_abstract_no_usage_relationship = response[\"choices\"][0][\"message\"][\"content\"]\n",
    "print(semi_automatic_related_work_gpt4_abstract_no_usage_relationship)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a458e73",
   "metadata": {},
   "outputs": [],
   "source": [
    "everything_collected = {\n",
    "    \"prefix\": prefix,\n",
    "    \"keyword\": keyword,\n",
    "    \"target_paper_name\": target_paper_name,\n",
    "    \"cited_jsons\": cited_jsons,\n",
    "    \"global_metadata\": global_metadata,\n",
    "    \"cited_facted_summaries\": cited_facted_summaries,\n",
    "    \"relationships_by_cited_paper\": relationships_by_cited_paper,\n",
    "    \"citation_spans_by_cited_paper\": citation_spans_by_cited_paper,\n",
    "    \"citation_intent_summary\": citation_intent_summary,\n",
    "    \"high_level_idea_prompt\": high_level_idea_prompt,\n",
    "    \"main_idea\": main_idea,\n",
    "    \"target_marker\": target_marker,\n",
    "    \"fully_automatic_generation_prompt\": fully_automatic_generation_prompt,\n",
    "    \"fully_automatic_related_work_gpt4\": fully_automatic_related_work_gpt4,\n",
    "    \"semi_automatic_generation_prompt\": semi_automatic_generation_prompt,\n",
    "    \"semi_automatic_related_work_gpt4\": semi_automatic_related_work_gpt4,\n",
    "    \"semi_automatic_generation_prompt_no_usage\": semi_automatic_generation_prompt_no_usage,\n",
    "    \"semi_automatic_related_work_gpt4_no_usage\": semi_automatic_related_work_gpt4_no_usage,\n",
    "    \"semi_automatic_generation_prompt_no_relationship\": semi_automatic_generation_prompt_no_relationship,\n",
    "    \"semi_automatic_related_work_gpt4_no_relationship\": semi_automatic_related_work_gpt4_no_relationship,\n",
    "    \"semi_automatic_generation_prompt_no_TAIC\": semi_automatic_generation_prompt_no_TAIC,\n",
    "    \"semi_automatic_related_work_gpt4_no_TAIC\": semi_automatic_related_work_gpt4_no_TAIC,\n",
    "    \"k\": k,\n",
    "    \"semi_automatic_generation_prompt_CTS\": semi_automatic_generation_prompt_CTS,\n",
    "    \"related_work_text_CTS\": related_work_text_CTS,\n",
    "    \"semi_automatic_generation_prompt_abstract\": semi_automatic_generation_prompt_abstract,\n",
    "    \"semi_automatic_related_work_gpt4_abstract\": semi_automatic_related_work_gpt4_abstract,\n",
    "    \"semi_automatic_generation_prompt_no_usage_relationship\": semi_automatic_generation_prompt_no_usage_relationship,\n",
    "    \"semi_automatic_related_work_gpt4_no_usage_relationship\": semi_automatic_related_work_gpt4_no_usage_relationship,\n",
    "    \"semi_automatic_generation_prompt_abstract_CTS\": semi_automatic_generation_prompt_abstract_CTS,\n",
    "    \"related_work_text_abstract_CTS\": related_work_text_abstract_CTS,\n",
    "    \"semi_automatic_generation_prompt_abstract_no_usage_relationship\": semi_automatic_generation_prompt_abstract_no_usage_relationship,\n",
    "    \"semi_automatic_related_work_gpt4_abstract_no_usage_relationship\": semi_automatic_related_work_gpt4_abstract_no_usage_relationship,\n",
    "    \"span_citation_mappings_lookup\": span_citation_mappings_lookup,\n",
    "}\n",
    "with open(prefix.replace(\"pdfs/\",\"dump_full.json\"),\"w\") as f:\n",
    "    json.dump(everything_collected, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddd6ed38",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}