{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "180f9bc1-03cc-4e31-babe-3f6c6ecb0167",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef6b0609-695f-4975-9970-f8b8350f953d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f986b468-c428-4ca3-9101-1cbabe6ad73f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from vllm import LLM, SamplingParams\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch.nn.functional as F\n",
    "import torch\n",
    "from transformers import AutoTokenizer\n",
    "from transformers import AutoModelForCausalLM\n",
    "import re\n",
    "import os\n",
    "#os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"1\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e0537c8-85cc-4bae-97de-6dd6f70ea5a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "llama = LLM(model='hugging-quants/Meta-Llama-3.1-70B-Instruct-AWQ-INT4', tensor_parallel_size = 2, download_dir = \"../meta_ai/\", gpu_memory_utilization=0.75, max_model_len=20000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0394142-a749-4995-abc8-bac884fea671",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_synthetic_histories(num_histories, llama_model):\n",
    "    np.random.seed(42)\n",
    "    tokenizer = llama_model.get_tokenizer()\n",
    "    prompts = []\n",
    "    cancer_types = np.random.choice(['breast', 'non-small cell lung', 'small cell lung', 'colorectal', 'pancreatic', 'urothelial', 'prostate', 'gastric', 'esophageal', 'thymoma', 'thymic carcinoma', 'adrenal', 'ovarian', 'endometrial', 'melanoma', 'renal cell', 'sarcoma', 'head and neck', 'Hodgkin lymphoma', 'Non-Hodgkin lymphoma', 'myeloma', 'acute myeloid leukemia', 'chronic myeloid leukemia', 'acute lymphoblastic leukemia', 'chronic lymphocytic leukemia/lymphoma', 'primary brain tumor'], size=num_histories) \n",
    "    splits = np.random.choice(['train','val','test'], size=num_histories, p=[0.8,0.1,0.1])\n",
    "    for i in range(num_histories):\n",
    "        messages = [\n",
    "            {'role':'system', 'content': \"\"\"Your job is to generate synthetic clinical histories for hypothetical patients with cancer.\n",
    "            You know all there is to know about cancer and its treatments, so be detailed.\n",
    "            The histories should be presented in chronological order as a sequence of events. Each event should begin with a date, and should then include some new development, such as a diagnosis, treatment, adverse event, progression, response to therapy, biomarker ascertainment, symptom burden, recurrence events, and so on.\n",
    "            \n",
    "        \"\"\"},      \n",
    "              \n",
    "            {'role':'user', 'content': \"\"\"Imagine a patient with cancer. \n",
    "            The cancer type is \"\"\" + cancer_types[i] + \"\"\".\n",
    "            Then, generate a very detailed synthetic clinical history for the patient. The patient might have any stage of disease. Use everything you know about cancer, including epidemiology, treatment options, outcomes, and heterogeneity in disease trajectories.\n",
    "            Do not mention transitions to hospice or death events.\n",
    "            Do not start with any demographics; just launch into the chronological history. Phrase it in the past tense. Dates should be in mm/dd/yyyy format. Output should be plain text, not Markdown.\n",
    "            The history should be approximately two pages long.\"\"\"}\n",
    "        ]\n",
    "    \n",
    "        prompts.append(tokenizer.apply_chat_template(conversation=messages, add_generation_prompt=True, tokenize=False))\n",
    "    \n",
    "\n",
    "    \n",
    "    responses = llama_model.generate(\n",
    "        prompts,   \n",
    "        SamplingParams(\n",
    "        temperature=1.0,\n",
    "        top_p=0.9,\n",
    "        max_tokens=5000,\n",
    "        stop_token_ids=[tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids(\"<|eot_id|>\")],  # KEYPOINT HERE\n",
    "    ))\n",
    "\n",
    "    response_texts = [x.outputs[0].text for x in responses]\n",
    "\n",
    "\n",
    "    return pd.DataFrame({'split':splits, 'cancer_type':cancer_types, 'patient_long_text':response_texts})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "109f2208-de29-43cf-a831-4609bdab225e",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = generate_synthetic_histories(30000, llama)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6119a0b4-5a63-4f91-a637-484b5e9dc29c",
   "metadata": {},
   "outputs": [],
   "source": [
    "results.to_csv('synthetic_histories_11-22-24.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71a24b19-5e1c-4c24-b13b-ac04c1e94bd2",
   "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}