{
 "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": "a8138f4e-6f45-4c98-b6b6-19370d53e7ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# llama = LLM(model='meta-llama/Meta-Llama-3.1-8B-Instruct', tensor_parallel_size = 2, \n",
    "#             gpu_memory_utilization=0.95,\n",
    "#             download_dir = \"../../\", max_model_len=120000)"
   ]
  },
  {
   "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.80, max_model_len=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0394142-a749-4995-abc8-bac884fea671",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_synthetic_imaging_reports(num_reports, llama_model):\n",
    "\n",
    "    tokenizer = llama_model.get_tokenizer()\n",
    "    prompts = []\n",
    "    scan_types = np.random.choice(['CT scan', 'MRI', 'Nuclear bone scan', 'PET-CT'], size=num_reports)\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_reports) \n",
    "\n",
    "    for i in range(num_reports):\n",
    "\n",
    "        messages = [\n",
    "            {'role':'system', 'content': \"\"\"Your job is to generate synthetic imaging reports for hypothetical patients with cancer.\n",
    "            You know all there is to know about cancer and its treatments, so be detailed.           \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 imaging report that might have been written about an imaging study performed for the patient. \n",
    "            The patient might have any stage of disease and be at any point along the disease trajectory. Use everything you know about cancer, including epidemiology, treatment, and heterogeneity in disease presentations.\n",
    "            The imaging study type is \"\"\" + scan_types[i] + \".\" + \"\"\"\n",
    "            The report should include a detailed \"Findings\" section followed by an \"Impression\" section.\n",
    "            The report should not include any treatment recommendations.\n",
    "            The imaging report should be approximately a full page 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=4000,\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({'cancer_type':cancer_types, 'scan_type':scan_types, 'synthetic_imaging_report':response_texts})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "109f2208-de29-43cf-a831-4609bdab225e",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = generate_synthetic_imaging_reports(10000, llama)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "443fa19c-8933-41d2-ba04-382902421e08",
   "metadata": {},
   "outputs": [],
   "source": [
    "results.synthetic_imaging_report.sample(n=1).iloc[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b40b06f0-69e9-48cc-9766-61d8d26178d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "results.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6119a0b4-5a63-4f91-a637-484b5e9dc29c",
   "metadata": {},
   "outputs": [],
   "source": [
    "results.to_csv('synthetic_imaging_reports.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.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}