Update README.md

README.md

@@ -1,29 +1,44 @@
 ---
 license: apache-2.0
-language:
-- en
-base_model:
-- mistralai/Mistral-7B-v0.1
-pipeline_tag: text2text-generation
 tags:
-- summerization
-- lora
-- medical
-- mistral
-- peft
-- mts-dialog
+  - medical
+  - summarization
+  - lora
+  - mistral
+  - dialogue
+  - peft
+model-index:
+  - name: Fine-tuned Mistral-7B (LoRA) on MTS-Dialog
+    results:
+      - task:
+          type: summarization
+        metrics:
+          - name: ROUGE-1
+            type: rouge
+            value: 0.1318
+          - name: ROUGE-2
+            type: rouge
+            value: 0.0456
+          - name: ROUGE-L
+            type: rouge
+            value: 0.0900
+          - name: BLEU
+            type: bleu
+            value: 0.0260
 ---
-# 🧠 LoRA Fine-Tuned Mistral-7B on MTS-Dialog Dataset
 
-This repository contains a LoRA fine-tuned version of the [`mistralai/Mistral-7B-v0.1`](https://huggingface.co/mistralai/Mistral-7B-v0.1) model on the **MTS-Dialog** dataset for medical dialogue summarization.
+# 🧠 LoRA Fine-Tuned Mistral-7B on MTS-Dialog
+
+This repository contains a LoRA fine-tuned version of [`mistralai/Mistral-7B-v0.1`](https://huggingface.co/mistralai/Mistral-7B-v0.1) for medical dialogue summarization, trained on the [MTS-Dialog](https://github.com/abachaa/MTS-Dialog) dataset.
 
 ---
 
 ## 📘 Model Summary
 
-- **Base Model**: Mistral-7B (v0.1)
-- **Technique**: LoRA (Low-Rank Adaptation)
-- **Framework**: 🤗 Hugging Face Transformers + PEFT + bitsandbytes
+- **Base Model**: `mistralai/Mistral-7B-v0.1`
+- **Fine-tuning Method**: LoRA (Low-Rank Adaptation)
+- **Frameworks**: 🤗 Transformers, PEFT, bitsandbytes
+- **Quantization**: 4-bit
 - **Task**: Medical dialogue summarization
 - **Dataset**: [MTS-Dialog](https://github.com/abachaa/MTS-Dialog)
 
@@ -31,65 +46,26 @@ This repository contains a LoRA fine-tuned version of the [`mistralai/Mistral-7B
 
 ## 🏥 Task Description
 
-The goal is to generate concise summaries of doctor-patient interactions, tailored to specific sections of the medical record (e.g., GENHX, HPI, ROS). This is useful for clinical documentation automation and decision support.
-
-Each training sample follows this pattern:
-
-"""
-Example 1:
-Dialogue:
-Doctor: Hello, Mrs. Smith. What seems to be troubling you today?
-Patient: I’ve been having shortness of breath and a mild cough for two weeks.
-Doctor: Any history of asthma or allergies?
-Patient: No, I’ve never had any breathing problems before.
-Summary:
-The patient, a middle-aged woman, presented with a two-week history of shortness of breath and mild cough without prior respiratory conditions. The physician asked about asthma/allergies, which the patient denied.
-
-Example 2:
-Dialogue:
-Doctor: Good morning. How are you feeling since your last visit?
-Patient: I still have a sharp pain in my right knee when I climb stairs.
-Doctor: Does the pain radiate anywhere else?
-Patient: No, it’s just in my knee. It started about a month ago.
-Summary:
-The patient continues to experience sharp knee pain exacerbated by stair climbing for one month, localized to the right knee with no radiation.
-"""
-
-new_dialogue_header = "GENHX"
-new_dialogue_text = """
-Doctor: What brings you back into the clinic today, miss?
-Patient: I've had chest pain for the last few days.
-Doctor: When did it start?
-"""
-
-inference_prompt = few_shot_prompts + f"""
-
-Now you:
-Summarize the following dialogue for section: {new_dialogue_header}
-{new_dialogue_text}
-Summary:
-"""
-
-
+This model is trained to summarize doctor-patient conversations into concise clinical notes, categorized by sections such as `GENHX`, `HPI`, `ROS`, etc. These summaries assist with EHR documentation and clinical decision-making.
 
 ---
 
-## 📊 Training Configuration
+## ⚙️ Training Configuration
 
-- **LoRA Rank**: `r=4`
-- **Epochs**: `3`
-- **Batch Size**: `4 (gradient_accumulation=4)`
-- **Learning Rate**: `3e-4`
-- **Quantization**: 4-bit using `bitsandbytes`
-- **Device**: `cuda:0` (single GPU)
+| Parameter          | Value         |
+|--------------------|---------------|
+| LoRA Rank          | 4             |
+| Epochs             | 3             |
+| Batch Size         | 4 (×4 grad. acc.) |
+| Learning Rate      | 3e-4          |
+| Device             | CUDA:0        |
+| Quantization       | 4-bit (bnb)   |
 
-Due to limited GPU resources (office laptop), training was intentionally short with a reduced LoRA rank and limited epochs. This led to **suboptimal performance**, which can be improved with longer training and higher-rank adapters.
+> ⚠️ Due to limited GPU resources (office laptop), training was constrained to 3 epochs and a small LoRA rank. Performance is expected to improve significantly with extended training and better hardware.
 
 ---
 
-## 📈 Evaluation Metrics
-
-Final validation metrics after 3 epochs:
+## 📊 Evaluation Metrics
 
 | Metric    | Score  |
 |-----------|--------|
@@ -98,55 +74,43 @@ Final validation metrics after 3 epochs:
 | ROUGE-L   | 0.0900 |
 | BLEU      | 0.0260 |
 
-> ⚠️ **Note**: These results are lower than expected due to low-rank LoRA (`r=4`) and only 3 epochs. Further tuning (e.g. `r=8`, `epochs=10`) on better GPUs will likely improve performance.
-
 ---
 
-## 💡 Prompting Examples
-
-**Input**:
+## 💡 Example Prompt
 
-
-# ---------------------------------------------
-# Example A (Influenza Suspect Dialogue)
-# ---------------------------------------------
-exampleA = """
-Doctor: Hello, Mr. Patel. Are you having any fever or chills?
-Patient: Yes, I’ve had a 102°F fever since yesterday and chills last night.
-Doctor: Any cough or stuffy nose?
-Patient: Mild cough and some congestion.
-Doctor: Do you have body aches?
-Patient: Yes, I feel sore all over.
+```text
+Summarize the following dialogue for section: GENHX
+Doctor: What brings you back into the clinic today, miss?
+Patient: I've had chest pain for the last few days.
+Doctor: When did it start?
 Summary:
-"""
-
-Example A Generated Summary:
- This is a case of influenza with fever, cough, and myalgia. The patient also has a history of asthma and hypertension. He has not been vaccinated against the flu this year.
-
-
----
 
-## 📁 Files
 
-- `config.json` – PEFT LoRA config
-- `adapter_model.bin` – LoRA adapter weights
-- `tokenizer/` – Tokenizer files
-- `README.md` – This model card
+## 🧪 Inference Code
 
----
-
-## 🔄 How to Use
-
-```python
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
 
 model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1", load_in_4bit=True)
 model = PeftModel.from_pretrained(model, "Imsachinsingh00/Fine_tuned_LoRA_Mistral_MTSDialog_Summarization")
+model.eval()
 
 tokenizer = AutoTokenizer.from_pretrained("Imsachinsingh00/Fine_tuned_LoRA_Mistral_MTSDialog_Summarization")
 
 prompt = "Summarize the following dialogue for section: HPI\nDoctor: Hello, what brings you in?\nPatient: I've been dizzy for two days.\nSummary:"
 inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 output = model.generate(**inputs, max_new_tokens=150)
-print(tokenizer.decode(output[0], skip_special_tokens=True))
+print(tokenizer.decode(output[0], skip_special_tokens=True))
+
+
+## 📁 Included Files
+
+- `config.json` – PEFT configuration for LoRA  
+- `adapter_model.bin` – LoRA adapter weights  
+- `tokenizer/` – Tokenizer files  
+- `README.md` – This model card  
+
+## 📌 Notes
+
+- 🚫 This is not a fully optimized clinical model — only a proof of concept.  
+- 💡 Consider training longer (`epochs=10`, `rank=8`) on GPUs with higher VRAM for better results.