Upload 3 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+medical_coding_train_1.csv filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,206 @@
+import pandas as pd
+import torch
+import os
+from transformers import AutoTokenizer, BitsAndBytesConfig, AutoModelForCausalLM
+from huggingface_hub import login
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.embeddings import SentenceTransformerEmbeddings
+from langchain.vectorstores import FAISS
+from langchain.chains import RetrievalQA
+from langchain.prompts import PromptTemplate
+from langchain.llms import HuggingFacePipeline
+from langchain_community.document_loaders.csv_loader import CSVLoader
+import transformers
+from langchain.schema import Document
+import gradio as gr
+import re
+
+
+model = "abnuel/MedGemma-4b-ICD"
+
+#tokenizer = AutoTokenizer.from_pretrained("abnuel/MedGemma-4b-ICD")
+
+SYSTEM_PROMPT = "You are an expert medical coder. Your task is to analyze the clinical description provided and output only the single, most appropriate ICD-10-CM code. Do not include any text, justification other than the code itself."
+
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+tokenizer = AutoTokenizer.from_pretrained(model)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+model = AutoModelForCausalLM.from_pretrained(
+    model,
+    quantization_config=bnb_config,
+    device_map="auto"
+)
+
+
+def generate_response(clinical_note):
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user", "content": f"Code the following: {clinical_note}"},
+    ]
+
+    # 3. Apply chat template and tokenize
+
+    inputs = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_dict=True,
+        return_tensors="pt"
+    ).to(model.device)
+
+    input_len = inputs["input_ids"].shape[-1]
+
+    # 4. Generate the response
+    with torch.inference_mode():
+        generation = model.generate(
+            **inputs,
+            max_new_tokens=200,      # Max length of the generated ICD codes
+            do_sample=False,         # Use greedy decoding for predictable output
+            temperature=0.0,         # Zero temperature for deterministic results
+        )
+
+    # 5. Decode the output
+    # Extract only the newly generated tokens
+    generation = generation[0][input_len:]
+    decoded_output = tokenizer.decode(generation, skip_special_tokens=True)
+
+    return decoded_output.strip()
+
+# --- Example Usage ---
+#test_note = "Sudden onset chest pain and shortness of breath. Initial diagnosis points towards unstable angina."
+
+
+#print(f"Clinical Note: {test_note}")
+#response = generate_response(test_note)
+#print(f"Generated ICD Codes: {response}")
+
+
+
+pipe = transformers.pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=50,
+    temperature=0.1,
+    do_sample=False,
+    pad_token_id=tokenizer.eos_token_id,
+
+)
+hf_llm = HuggingFacePipeline(pipeline=pipe)
+
+
+df = pd.read_csv("medical_coding_train_1.csv")
+
+documents = [
+    Document(
+        page_content=f"note: {row['note']}\nicd_code: {row['icd_codes']}",
+        metadata={"icd_code": row["icd_codes"]}
+    )
+    for _, row in df.iterrows()
+]
+
+
+# 2. Chunk Documents
+text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+docs = text_splitter.split_documents(documents)
+
+# 3. Create Embeddings and Vector Store (FAISS)
+embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
+db = FAISS.from_documents(docs, embeddings)
+retriever = db.as_retriever(search_kwargs={"k": 2}) 
+
+
+RAG_PROMPT_TEMPLATE = """
+You are an expert medical coder.
+Your task is to determine the most accurate ICD-10-CM code for the given clinical note.
+
+Use ONLY the following context (which may include ICD codes from similar cases).
+If you cannot determine a match from the provided context, respond exactly with:
+"I cannot find the code in the provided documents."
+
+Return ONLY the ICD-10-CM code itself — no explanation, no text, no punctuation.
+
+Context:
+{context}
+
+Clinical Note:
+{question}
+
+ICD-10-CM Code:
+"""
+
+
+rag_prompt = PromptTemplate.from_template(RAG_PROMPT_TEMPLATE)
+#direct_chain = LLMChain(llm=hf_llm, prompt=rag_prompt)
+
+
+# 5. Create the QA Chain
+qa_chain = RetrievalQA.from_chain_type(
+    llm=hf_llm,
+    chain_type="stuff",
+    retriever=retriever,
+    return_source_documents=False,
+    chain_type_kwargs={"prompt": rag_prompt}
+)
+
+def extract_icd_code(text):
+    # Pattern to match "ICD-10-CM Code:" followed by the code
+    pattern = r'ICD-10-CM Code:\s*([A-Z0-9.]+)'
+
+    match = re.search(pattern, text)
+
+    if match:
+        return match.group(1)
+    return None
+
+def generate_code_rag(clinical_note, retriever, threshold=0.35):
+    """Generates the ICD code using RAG."""
+
+    # Format the user question for the RAG prompt template
+    query = f"Code the following: {clinical_note}"
+
+     # Step 1: Retrieve docs
+    docs_and_scores = db.similarity_search_with_score(query, k=2)
+
+    # Step 2: Filter by similarity threshold
+    relevant_docs = [doc for doc, score in docs_and_scores if score > threshold]
+
+    
+
+    if relevant_docs:
+        #print(qa_chain)
+        result = qa_chain({"query": query})["result"]
+        
+        #answer = result['result']
+        icd_code = extract_icd_code(result)
+        #print(icd_code)
+        if icd_code == None:
+          print("I got here")
+          result = generate_response(clinical_note)
+          return result
+        else:
+          return icd_code
+    
+        # Step 4: Otherwise, use LLM directly (no context)
+        
+
+# Create the Gradio Interface
+gr.Interface(
+    fn=generate_code_rag,
+    inputs=gr.Textbox(lines=5, label="Enter Clinical Note Here", placeholder="e.g., Patient presented with simple laceration of the left hand."),
+    outputs=gr.Textbox(label="Predicted ICD-10 Code"),
+    title="ClaimSwift Medical Coding",
+    description="",
+    examples=[
+        ["Benign neoplasm of peripheral nerves and autonomic nervous system of face, head, and neck"],
+        ["Sudden onset chest pain and shortness of breath. Initial diagnosis points towards unstable angina."],
+        ["Simple laceration of the left hand without foreign body."],
+    ]
+).launch(server_name="0.0.0.0", server_port=7860)

medical_coding_train_1.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a812e716b9a718799315f93c2aa4ad7efe0c9b93141cdca94f110aa29653af9f
+size 15109530

requirements.txt

@@ -0,0 +1,12 @@
+pandas
+langchain
+langchain-community
+langchain-core
+transformers
+torch
+accelerate
+bitsandbytes
+sentence-transformers
+faiss-cpu
+gradio
+peft