workspace.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "afba65fd",
   "metadata": {},
   "source": [
    "- ubuntu 18.04\n",
    "- cuda 12.1 driver 530.41\n",
    "- 2x 3090\n",
    "- AMD Ryzen Threadripper 1920X 12-Core\n",
    "- 128 gb Ram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6276471a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import time\n",
    "import nemo.collections.asr as nemo_asr\n",
    "\n",
    "from src.utils import extract_soap_sections\n",
    "from llama_cpp import Llama\n",
    "from pydub import AudioSegment\n",
    "from transformers import pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b5846dee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2024-06-05 09:23:29 cloud:58] Found existing object /root/.cache/torch/NeMo/NeMo_1.23.0/QuartzNet15x5Base-En/2b066be39e9294d7100fb176ec817722/QuartzNet15x5Base-En.nemo.\n",
      "[NeMo I 2024-06-05 09:23:29 cloud:64] Re-using file from: /root/.cache/torch/NeMo/NeMo_1.23.0/QuartzNet15x5Base-En/2b066be39e9294d7100fb176ec817722/QuartzNet15x5Base-En.nemo\n",
      "[NeMo I 2024-06-05 09:23:29 common:924] Instantiating model from pre-trained checkpoint\n",
      "[NeMo I 2024-06-05 09:23:30 features:289] PADDING: 16\n",
      "[NeMo I 2024-06-05 09:23:31 save_restore_connector:249] Model EncDecCTCModel was successfully restored from /root/.cache/torch/NeMo/NeMo_1.23.0/QuartzNet15x5Base-En/2b066be39e9294d7100fb176ec817722/QuartzNet15x5Base-En.nemo.\n"
     ]
    }
   ],
   "source": [
    "quartznet = nemo_asr.models.EncDecCTCModel.from_pretrained(\n",
    "    model_name=\"QuartzNet15x5Base-En\",\n",
    "    map_location='cuda'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "88f16917",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "llama_model_loader: loaded meta data with 19 key-value pairs and 291 tensors from llama-2-7b-chat.Q4_K_M.gguf (version GGUF V2)\n",
      "llama_model_loader: Dumping metadata keys/values. Note: KV overrides do not apply in this output.\n",
      "llama_model_loader: - kv   0:                       general.architecture str              = llama\n",
      "llama_model_loader: - kv   1:                               general.name str              = LLaMA v2\n",
      "llama_model_loader: - kv   2:                       llama.context_length u32              = 4096\n",
      "llama_model_loader: - kv   3:                     llama.embedding_length u32              = 4096\n",
      "llama_model_loader: - kv   4:                          llama.block_count u32              = 32\n",
      "llama_model_loader: - kv   5:                  llama.feed_forward_length u32              = 11008\n",
      "llama_model_loader: - kv   6:                 llama.rope.dimension_count u32              = 128\n",
      "llama_model_loader: - kv   7:                 llama.attention.head_count u32              = 32\n",
      "llama_model_loader: - kv   8:              llama.attention.head_count_kv u32              = 32\n",
      "llama_model_loader: - kv   9:     llama.attention.layer_norm_rms_epsilon f32              = 0.000001\n",
      "llama_model_loader: - kv  10:                          general.file_type u32              = 15\n",
      "llama_model_loader: - kv  11:                       tokenizer.ggml.model str              = llama\n",
      "llama_model_loader: - kv  12:                      tokenizer.ggml.tokens arr[str,32000]   = [\"<unk>\", \"<s>\", \"</s>\", \"<0x00>\", \"<...\n",
      "llama_model_loader: - kv  13:                      tokenizer.ggml.scores arr[f32,32000]   = [0.000000, 0.000000, 0.000000, 0.0000...\n",
      "llama_model_loader: - kv  14:                  tokenizer.ggml.token_type arr[i32,32000]   = [2, 3, 3, 6, 6, 6, 6, 6, 6, 6, 6, 6, ...\n",
      "llama_model_loader: - kv  15:                tokenizer.ggml.bos_token_id u32              = 1\n",
      "llama_model_loader: - kv  16:                tokenizer.ggml.eos_token_id u32              = 2\n",
      "llama_model_loader: - kv  17:            tokenizer.ggml.unknown_token_id u32              = 0\n",
      "llama_model_loader: - kv  18:               general.quantization_version u32              = 2\n",
      "llama_model_loader: - type  f32:   65 tensors\n",
      "llama_model_loader: - type q4_K:  193 tensors\n",
      "llama_model_loader: - type q6_K:   33 tensors\n",
      "llm_load_vocab: special tokens cache size = 259\n",
      "llm_load_vocab: token to piece cache size = 0.1684 MB\n",
      "llm_load_print_meta: format           = GGUF V2\n",
      "llm_load_print_meta: arch             = llama\n",
      "llm_load_print_meta: vocab type       = SPM\n",
      "llm_load_print_meta: n_vocab          = 32000\n",
      "llm_load_print_meta: n_merges         = 0\n",
      "llm_load_print_meta: n_ctx_train      = 4096\n",
      "llm_load_print_meta: n_embd           = 4096\n",
      "llm_load_print_meta: n_head           = 32\n",
      "llm_load_print_meta: n_head_kv        = 32\n",
      "llm_load_print_meta: n_layer          = 32\n",
      "llm_load_print_meta: n_rot            = 128\n",
      "llm_load_print_meta: n_embd_head_k    = 128\n",
      "llm_load_print_meta: n_embd_head_v    = 128\n",
      "llm_load_print_meta: n_gqa            = 1\n",
      "llm_load_print_meta: n_embd_k_gqa     = 4096\n",
      "llm_load_print_meta: n_embd_v_gqa     = 4096\n",
      "llm_load_print_meta: f_norm_eps       = 0.0e+00\n",
      "llm_load_print_meta: f_norm_rms_eps   = 1.0e-06\n",
      "llm_load_print_meta: f_clamp_kqv      = 0.0e+00\n",
      "llm_load_print_meta: f_max_alibi_bias = 0.0e+00\n",
      "llm_load_print_meta: f_logit_scale    = 0.0e+00\n",
      "llm_load_print_meta: n_ff             = 11008\n",
      "llm_load_print_meta: n_expert         = 0\n",
      "llm_load_print_meta: n_expert_used    = 0\n",
      "llm_load_print_meta: causal attn      = 1\n",
      "llm_load_print_meta: pooling type     = 0\n",
      "llm_load_print_meta: rope type        = 0\n",
      "llm_load_print_meta: rope scaling     = linear\n",
      "llm_load_print_meta: freq_base_train  = 10000.0\n",
      "llm_load_print_meta: freq_scale_train = 1\n",
      "llm_load_print_meta: n_yarn_orig_ctx  = 4096\n",
      "llm_load_print_meta: rope_finetuned   = unknown\n",
      "llm_load_print_meta: ssm_d_conv       = 0\n",
      "llm_load_print_meta: ssm_d_inner      = 0\n",
      "llm_load_print_meta: ssm_d_state      = 0\n",
      "llm_load_print_meta: ssm_dt_rank      = 0\n",
      "llm_load_print_meta: model type       = 7B\n",
      "llm_load_print_meta: model ftype      = Q4_K - Medium\n",
      "llm_load_print_meta: model params     = 6.74 B\n",
      "llm_load_print_meta: model size       = 3.80 GiB (4.84 BPW) \n",
      "llm_load_print_meta: general.name     = LLaMA v2\n",
      "llm_load_print_meta: BOS token        = 1 '<s>'\n",
      "llm_load_print_meta: EOS token        = 2 '</s>'\n",
      "llm_load_print_meta: UNK token        = 0 '<unk>'\n",
      "llm_load_print_meta: LF token         = 13 '<0x0A>'\n",
      "llm_load_tensors: ggml ctx size =    0.15 MiB\n",
      "llm_load_tensors:        CPU buffer size =  3891.24 MiB\n",
      "..................................................................................................\n",
      "llama_new_context_with_model: n_ctx      = 4096\n",
      "llama_new_context_with_model: n_batch    = 512\n",
      "llama_new_context_with_model: n_ubatch   = 512\n",
      "llama_new_context_with_model: flash_attn = 0\n",
      "llama_new_context_with_model: freq_base  = 10000.0\n",
      "llama_new_context_with_model: freq_scale = 1\n",
      "llama_kv_cache_init:        CPU KV buffer size =  2048.00 MiB\n",
      "llama_new_context_with_model: KV self size  = 2048.00 MiB, K (f16): 1024.00 MiB, V (f16): 1024.00 MiB\n",
      "llama_new_context_with_model:        CPU  output buffer size =     0.12 MiB\n",
      "llama_new_context_with_model:        CPU compute buffer size =   296.01 MiB\n",
      "llama_new_context_with_model: graph nodes  = 1030\n",
      "llama_new_context_with_model: graph splits = 1\n",
      "AVX = 1 | AVX_VNNI = 0 | AVX2 = 1 | AVX512 = 0 | AVX512_VBMI = 0 | AVX512_VNNI = 0 | AVX512_BF16 = 0 | FMA = 1 | NEON = 0 | SVE = 0 | ARM_FMA = 0 | F16C = 1 | FP16_VA = 0 | WASM_SIMD = 0 | BLAS = 0 | SSE3 = 1 | SSSE3 = 1 | VSX = 0 | MATMUL_INT8 = 0 | LLAMAFILE = 1 | \n",
      "Model metadata: {'tokenizer.ggml.unknown_token_id': '0', 'tokenizer.ggml.eos_token_id': '2', 'general.architecture': 'llama', 'llama.context_length': '4096', 'general.name': 'LLaMA v2', 'llama.embedding_length': '4096', 'llama.feed_forward_length': '11008', 'llama.attention.layer_norm_rms_epsilon': '0.000001', 'llama.rope.dimension_count': '128', 'llama.attention.head_count': '32', 'tokenizer.ggml.bos_token_id': '1', 'llama.block_count': '32', 'llama.attention.head_count_kv': '32', 'general.quantization_version': '2', 'tokenizer.ggml.model': 'llama', 'general.file_type': '15'}\n",
      "Using fallback chat format: llama-2\n"
     ]
    }
   ],
   "source": [
    "lcpp_llm = Llama(\n",
    "    model_path='llama-2-7b-chat.Q4_K_M.gguf',\n",
    "    n_ctx=4096,\n",
    "    n_threads=24,) # CPU cores\n",
    "#     n_batch=512, # Should be between 1 and n_ctx, consider the amount of VRAM in your GPU.\n",
    "#     n_gpu_layers=32 # Change this value based on your model and your GPU VRAM pool.\n",
    "#     )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "899258d5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Logging before flag parsing goes to stderr.\n",
      "W0605 09:23:39.196702 140267760935360 modeling_phi3.py:62] `flash-attention` package not found, consider installing for better performance: No module named 'flash_attn'.\n",
      "W0605 09:23:39.197815 140267760935360 modeling_phi3.py:66] Current `flash-attention` does not support `window_size`. Either upgrade or use `attn_implementation='eager'`.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bc787d2cd37649dcbb871b0f15f7c8e5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You are using the default legacy behaviour of the <class 'transformers.models.llama.tokenization_llama_fast.LlamaTokenizerFast'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "pipe = pipeline(\n",
    "    \"text-generation\", \n",
    "    model=\"omi-health/sum-small\", \n",
    "    trust_remote_code=True,\n",
    "    max_length=2100,\n",
    "    device='cuda'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "954778db",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Mb move to utils\n",
    "def pipeline_processing(path_to_wav: str, llm: bool = True):\n",
    "    promt = (\n",
    "        \"You are an expert medical professor assisting in the creation of medically accurate SOAP summaries. Please ensure the response only follows the structured format: S:, O:, A:, P: without using markdown or special formatting and without your text. Avoid informal or conversational language in the all sections.\",\n",
    "    )\n",
    "\n",
    "    ### Validate wav, resample, converting\n",
    "    audio = AudioSegment.from_file(path_to_wav)\n",
    "    if audio.channels != (1):\n",
    "        audio = audio.set_channels(1)\n",
    "    if audio.frame_rate != 16000:\n",
    "        audio = audio.set_frame_rate(16000)\n",
    "    audio.export(\"conv_med.wav\", format=\"wav\")\n",
    "\n",
    "    # ASR\n",
    "    text_from_asr = quartznet.transcribe([\"conv_med.wav\"])[0]\n",
    "    soap_request = [\n",
    "        {\n",
    "            \"role\": \"system\",\n",
    "            \"content\": promt,\n",
    "        },\n",
    "        {\"role\": \"user\", \"content\": text_from_asr},\n",
    "    ]\n",
    "    if llm:\n",
    "        outputs_from_model = lcpp_llm.create_chat_completion(soap_request)[\"choices\"][0][\"message\"][\"content\"]\n",
    "        results_soap = extract_soap_sections(outputs_from_model)\n",
    "    else:\n",
    "        outputs_from_model = pipe(soap_request)[0][\"generated_text\"][2][\"content\"]\n",
    "        results_soap = extract_soap_sections(outputs_from_model)\n",
    "\n",
    "    return text_from_asr, promt, results_soap, outputs_from_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "817ede91",
   "metadata": {},
   "outputs": [],
   "source": [
    "wav_path = '/path/to/wav'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0246b74",
   "metadata": {},
   "source": [
    "You can use %%time or %%timeit\n",
    "```python\n",
    "%%time\n",
    "start = time.time()\n",
    "text_from_asr, results, outputs = pipeline_processing('conversation_part_1.mp3', llm=False)\n",
    "end = time.time() - start\n",
    "print(end)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4e635157",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3571430ccdfd4948a72a065f558a7e1d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Truncation was not explicitly activated but `max_length` is provided a specific value, please use `truncation=True` to explicitly truncate examples to max length. Defaulting to 'longest_first' truncation strategy. If you encode pairs of sequences (GLUE-style) with the tokenizer you can select this strategy more precisely by providing a specific strategy to `truncation`.\n",
      "W0605 09:24:32.937731 140267760935360 logging.py:328] You are not running the flash-attention implementation, expect numerical differences.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "11.693514108657837\n"
     ]
    }
   ],
   "source": [
    "start = time.time()\n",
    "text_from_asr, results, outputs = pipeline_processing(wav_path, llm=False)\n",
    "end = time.time() - start\n",
    "print(end)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "980802ca",
   "metadata": {},
   "source": [
    "### Llama"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "8a339406",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "88f457d3c9b3450bbe186f4dc5dc4213",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "llama_tokenize_internal: Added a BOS token to the prompt as specified by the model but the prompt also starts with a BOS token. So now the final prompt starts with 2 BOS tokens. Are you sure this is what you want?\n",
      "\n",
      "llama_print_timings:        load time =   15773.86 ms\n",
      "llama_print_timings:      sample time =      53.46 ms /   158 runs   (    0.34 ms per token,  2955.37 tokens per second)\n",
      "llama_print_timings: prompt eval time =   15773.14 ms /   441 tokens (   35.77 ms per token,    27.96 tokens per second)\n",
      "llama_print_timings:        eval time =   14002.74 ms /   157 runs   (   89.19 ms per token,    11.21 tokens per second)\n",
      "llama_print_timings:       total time =   29929.29 ms /   598 tokens\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "30.41976833343506\n"
     ]
    }
   ],
   "source": [
    "start = time.time()\n",
    "text_from_asr, results, outputs = pipeline_processing(wav_path, llm=True)\n",
    "end = time.time() - start\n",
    "print(end)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1d72bc0",
   "metadata": {},
   "source": [
    "### Pyannote"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "95eb1dc3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<pyannote.audio.pipelines.speaker_diarization.SpeakerDiarization at 0x7f6168eae9e0>"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pyannote.audio import Pipeline\n",
    "import torch\n",
    "\n",
    "# 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\n",
    "\n",
    "pipeline = Pipeline.from_pretrained(\n",
    "  \"pyannote/speaker-diarization-3.1\",\n",
    "    use_auth_token=\"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\")\n",
    "\n",
    "pipeline.to(torch.device('cuda'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "1243f3f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "diarization = pipeline(\"\", min_speakers=1, max_speakers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "05e66eae",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyannote.core import Segment, Annotation\n",
    "\n",
    "def merge_segments(diarization):\n",
    "    merged_segments = []\n",
    "    current_segment, current_speaker = None, None\n",
    "\n",
    "    for segment, _, speaker in diarization.itertracks(yield_label=True):\n",
    "        if speaker == current_speaker:\n",
    "            current_segment = Segment(current_segment.start, segment.end)\n",
    "        else:\n",
    "            if current_segment is not None:\n",
    "                merged_segments.append((current_segment, current_speaker))\n",
    "            current_segment, current_speaker = segment, speaker\n",
    "    \n",
    "    if current_segment is not None:\n",
    "        merged_segments.append((current_segment, current_speaker))\n",
    "    \n",
    "    return merged_segments\n",
    "\n",
    "# Объединение сегментов\n",
    "merged_segments = merge_segments(diarization)\n",
    "\n",
    "audio = AudioSegment.from_file(\"conv1.wav\")\n",
    "if audio.channels != 1:\n",
    "    audio = audio.set_channels(1)\n",
    "if audio.frame_rate != 16000:\n",
    "    audio = audio.set_frame_rate(16000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "521f5c38",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "eed9f74e1e2142a1bf71cc1868afd52f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "71cca63abcc745f08cf404cc49c79e4d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b0587f80fd4f44ea8f6b43de6cace4a3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b6e995cdd88c4c60a5ca26ea00e88e79",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4fa88c57e16b47c0b0b34dcaefb50307",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "355cbd93730a4e28ba4f51ebb28702f1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "70f1a957540e4362a7320d94ffbb84a0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4ccf9f9ee1e14cf99c6493ac6136c1f3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cc53cf407d004d0b84a07ba14b49f66a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "598588a9be90429dbd84b8e8f3fc84e0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bf992046892b4c01917b068c89ecd165",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5d92e3e362ff4a5e8d5657b7168a67ee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7bc6ef5c1f1c4c36abc293a4258355cd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "433a0bab60314d439e4cfd7bb5fbe986",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7c438be5f70541bcab02eca87aa55695",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "41f848fe879341d8bccdb544e6bcc705",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "51bb592b48f94c4cb381e14861c2d7e5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d253987b047f4aac9dc17ed3f3ba37d7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6f44ea13111149b3871cf93aea22b710",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import io\n",
    "import numpy as np\n",
    "import soundfile as sf\n",
    "\n",
    "results = []\n",
    "\n",
    "for segment, speaker in merged_segments:\n",
    "    start_ms = int(segment.start * 1000)  # перевод в миллисекунды\n",
    "    end_ms = int(segment.end * 1000)  # перевод в миллисекунды\n",
    "    segment_audio = audio[start_ms:end_ms]\n",
    "    segment_audio.export('segment.wav', format='wav')\n",
    "    text_from_asr = quartznet.transcribe(['segment.wav'])[0]\n",
    "    results.append((speaker, text_from_asr))\n",
    "\n",
    "    \n",
    "promt = 'You are an expert medical professor assisting in the creation of medically accurate SOAP summaries. Please ensure the response only follows the structured format: S:, O:, A:, P: without using markdown or special formatting and without [Sure, here is the SOAP note for the dental visit:]. Avoid informal or conversational language in the all sections.',\n",
    "soap_request = [\n",
    "    {\n",
    "        'role': 'system',\n",
    "        'content': promt,\n",
    "    }\n",
    "]\n",
    "\n",
    "role_map = {\n",
    "    'SPEAKER_00': 'user',\n",
    "    'SPEAKER_01': 'doctor'\n",
    "}\n",
    "\n",
    "for speaker, text in results:\n",
    "    soap_request.append({\n",
    "        'role': role_map.get(speaker, 'unknown'),\n",
    "        'content': text\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "c1035161",
   "metadata": {},
   "outputs": [],
   "source": [
    "promt = 'You are an expert medical professor assisting in the creation of medically accurate SOAP summaries. Please ensure the response only follows the structured format: S:, O:, A:, P: without using markdown or special formatting and without your text. Avoid informal or conversational language in the all sections.',\n",
    "def pipeline_processing(wav: str, llm: bool = True):\n",
    "    path_to_wav = wav\n",
    "    \n",
    "    audio = AudioSegment.from_file(path_to_wav)\n",
    "    if audio.channels != (1):\n",
    "        audio = audio.set_channels(1)\n",
    "    if audio.frame_rate != 16000:\n",
    "        audio = audio.set_frame_rate(16000)\n",
    "    audio.export('conv_med.wav', format='wav')\n",
    "\n",
    "    diarization = pipeline('conv_med.wav', min_speakers=1, max_speakers=2)\n",
    "    \n",
    "    merged_segments = merge_segments(diarization)\n",
    "\n",
    "    results = []\n",
    "    for segment, speaker in merged_segments:\n",
    "        start_ms = int(segment.start * 1000)\n",
    "        end_ms = int(segment.end * 1000)\n",
    "        segment_audio = audio[start_ms:end_ms]\n",
    "        segment_audio.export('segment.wav', format='wav')\n",
    "        text_from_asr = quartznet.transcribe(['segment.wav'])[0]\n",
    "        results.append((speaker, text_from_asr))\n",
    "        \n",
    "    soap_request = [\n",
    "        {\n",
    "            'role': 'system',\n",
    "            'content': promt,\n",
    "        }\n",
    "    ]\n",
    "\n",
    "    role_map = {\n",
    "        'SPEAKER_00': 'user',\n",
    "        'SPEAKER_01': 'doctor'\n",
    "    }\n",
    "\n",
    "    for speaker, text in results:\n",
    "        soap_request.append({\n",
    "            'role': role_map.get(speaker, 'unknown'),\n",
    "            'content': text\n",
    "        })\n",
    "        \n",
    "    if llm:\n",
    "        outputs = lcpp_llm.create_chat_completion(soap_request)['choices'][0]['message']['content']\n",
    "        results_soap = extract_soap_sections(outputs)\n",
    "    else:\n",
    "        outputs = pipe(soap_request)\n",
    "        results_soap = extract_soap_sections(outputs[0]['generated_text'][2]['content'])\n",
    "    \n",
    "    return text_from_asr, results_soap, outputs, soap_request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "id": "5b175a35",
   "metadata": {
    "collapsed": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "41cb9466eebe48a4ac5b408a82004bc6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e3bc2226fe5a48c7ba518069b5696270",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a9da4490500b4e0b9c9356fcd2e11312",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ae634254602a45c59f3d2dc6b2c603df",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3fd60af7d8c646409831cb9b9d694792",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cc218ccccae549ec95d34109ebbed60d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "eba11916c2894af9aabff11bb889e93b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "34c997b3288446d9bcc4f5b69303c9c5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3624c9fdef084c4696520e39d02fa1dc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d39880376d9f488185ea15a9ae1f3055",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "57433185f3324526924903c506e7dc60",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b821cdd832b84f438f34942cd898fc43",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "df990b389e414ab9879849996f2ff6d8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2d39508beace40e198a1b24ae4a0d581",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6985bf67acb94227bb0dd83a5a5d95e2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5fa213f116c944fe9245b9bd63d5346b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7297ded789734289bb854b2422cfc8ec",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "595eab6466e04246901d978aafbcc275",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "de9a37ea39d449a2bbaf8aeab27dcc03",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7d553ba592d947daa18f8336a3b7a112",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8f7cbbc1f02248b5b36ac36021c3b47b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Transcribing:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "llama_tokenize_internal: Added a BOS token to the prompt as specified by the model but the prompt also starts with a BOS token. So now the final prompt starts with 2 BOS tokens. Are you sure this is what you want?\n",
      "Llama.generate: prefix-match hit\n",
      "\n",
      "llama_print_timings:        load time =   12913.88 ms\n",
      "llama_print_timings:      sample time =      20.01 ms /    61 runs   (    0.33 ms per token,  3048.32 tokens per second)\n",
      "llama_print_timings: prompt eval time =   19243.46 ms /   531 tokens (   36.24 ms per token,    27.59 tokens per second)\n",
      "llama_print_timings:        eval time =    5440.24 ms /    60 runs   (   90.67 ms per token,    11.03 tokens per second)\n",
      "llama_print_timings:       total time =   24736.58 ms /   591 tokens\n"
     ]
    }
   ],
   "source": [
    "start = time.time()\n",
    "text_from_asr, results, outputs, soap_request = pipeline_processing('conversation_2.m4a')\n",
    "end = time.time() - start"
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}