Upload LLM2Vec4CXR fine-tuned model

README.md

@@ -1,60 +1,194 @@
-# llm2vec4cxr
-
 ---
 license: mit
+base_model: microsoft/LLM2CLIP-Llama-3.2-1B-Instruct-CC-Finetuned
+tags:
+- text-embeddings
+- sentence-transformers
+- llm2vec
+- medical
+- chest-xray
+- radiology
+- clinical-nlp
+language:
+- en
+pipeline_tag: feature-extraction
+library_name: transformers
 ---
 
+# LLM2Vec4CXR - Fine-tuned Model for Chest X-ray Report Analysis
+
+This model is a fine-tuned version of [microsoft/LLM2CLIP-Llama-3.2-1B-Instruct-CC-Finetuned](https://huggingface.co/microsoft/LLM2CLIP-Llama-3.2-1B-Instruct-CC-Finetuned) specifically optimized for chest X-ray report analysis and medical text understanding.
+
+## Model Description
+
+LLM2Vec4CXR is a bidirectional language model that converts the base decoder-only LLM into a text encoder optimized for medical text embeddings. The model has been fully fine-tuned with modified pooling strategy (`latent_attention`) to better capture semantic relationships in chest X-ray reports.
+
+### Key Features
+
+- **Base Architecture**: LLM2CLIP-Llama-3.2-1B-Instruct
+- **Pooling Mode**: Latent Attention (modified from original)
+- **Bidirectional Processing**: Enabled for better context understanding
+- **Medical Domain**: Specialized for chest X-ray report analysis
+- **Max Length**: 512 tokens
+- **Precision**: bfloat16
+
+## Training Details
+
+### Training Data
+- Fully fine-tuned on chest X-ray reports and medical text data
+- Training focused on understanding pleural effusion status and other chest X-ray findings
+
+### Training Configuration
+- **Pooling Mode**: `latent_attention` (modified from base model)
+- **Enable Bidirectional**: True
+- **Max Length**: 512
+- **Torch Dtype**: bfloat16
+- **Full Fine-tuning**: All model weights were updated during training
+
 ## Usage
 
-### Pytorch version
-pip install llm2vec
+### Installation
 
-```python
-import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+```bash
+pip install torch transformers
+# Also requires the LLM2Vec wrapper - see the original repository for installation
+```
+
+### Basic Usage
 
+```python
 import torch
-from torch import nn
-from llm2vec import LLM2Vec
-from transformers import AutoTokenizer, AutoModel, AutoConfig
-
-class LLM2VecWrapper(LLM2Vec):
-    def prepare_for_tokenization(self, text):
-        text = (
-            "<|start_header_id|>user<|end_header_id|>\n\n"
-            + text.strip()
-            + "<|eot_id|>"
-        )
-        return  text
+import torch.nn.functional as F
+from llm2vec_wrapper import LLM2VecWrapper as LLM2Vec
+
+# Load the model
+model = LLM2Vec.from_pretrained(
+    base_model_name_or_path='lukeingawesome/llm2vec4cxr',
+    enable_bidirectional=True,
+    pooling_mode="latent_attention",
+    max_length=512,
+    torch_dtype=torch.bfloat16,
+)
+
+# Configure tokenizer
+tokenizer = model.tokenizer
+tokenizer.padding_side = 'left'
+
+# Example usage for chest X-ray report analysis
+def encode_text(text):
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
+    with torch.no_grad():
+        embeddings = model(inputs)
+    return embeddings
+
+# Example with medical text
+report = "There is a small increase in the left-sided effusion. There continues to be volume loss at both bases."
+embedding = encode_text(report)
+```
+
+### Advanced Usage with Separator-based Processing
+
+The model supports special separator-based processing for instruction-following tasks:
+
+```python
+def tokenize_with_separator(texts, tokenizer, max_length):
+    """Tokenize texts with special handling for separator-based splitting."""
+    texts_2 = []
+    original_texts = []
+    separator = '!@#$%^&*()'
+    
+    for text in texts:
+        parts = text.split(separator)
+        texts_2.append(parts[1] if len(parts) > 1 else "")
+        original_texts.append("".join(parts))
+
+    tokenized = tokenizer(
+        original_texts,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=max_length,
+    )
     
-class LlamaVec_1B_FeatureExtractor(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-        model_path = 'lukeingawesome/llm2vec4cxr'
-        config = AutoConfig.from_pretrained(model_path)
-
-        model = AutoModel.from_pretrained(model_path, config=config, trust_remote_code=True, torch_dtype=torch.bfloat16)
-        tokenizer = AutoTokenizer.from_pretrained(model_path)
-
-        tokenizer.pad_token = tokenizer.eos_token
-        tokenizer.padding_side = "left"
-
-        self.l2v = LLM2VecWrapper(model, tokenizer, pooling_mode="mean", max_length=512, skip_instruction=True)
-        
-    def extract_features(self, text):
-        with torch.amp.autocast('cuda'):
-            reps_norm = self.l2v.encode(text)
-            reps_norm = torch.nn.functional.normalize(reps_norm, p=2, dim=1)
-        return {"embeds": reps_norm}
-
-text_model = LlamaVec_1B_FeatureExtractor()
-captions = ["Pleural effusion is noted in the right lower lobe. No pneumothorax. No consolidation."]
-embeddings = text_model.extract_features(captions)
+    # Create embedding masks for the separated parts
+    embed_mask = None
+    for t_i, t in enumerate(texts_2):
+        ids = tokenizer([t], return_tensors="pt", padding=True, truncation=True, 
+                       max_length=max_length, add_special_tokens=False)
+        e_m = torch.zeros_like(tokenized["attention_mask"][t_i])
+        if len(ids["input_ids"][0]) > 0:
+            e_m[-len(ids["input_ids"][0]):] = torch.ones(len(ids["input_ids"][0]))
+        if embed_mask is None:
+            embed_mask = e_m.unsqueeze(0)
+        else:
+            embed_mask = torch.cat((embed_mask, e_m.unsqueeze(0)), dim=0)
+
+    tokenized["embed_mask"] = embed_mask
+    return tokenized
+
+# Example with instruction and report
+separator = '!@#$%^&*()'
+instruction = 'Determine the change or the status of the pleural effusion.'
+report = 'There is a small increase in the left-sided effusion.'
+text = instruction + separator + report
+
+tokenized = tokenize_with_separator([text], tokenizer, 512)
+embedding = model(tokenized)
 ```
 
-## BibTeX & Citation
+## Evaluation
+
+The model has been evaluated on chest X-ray report analysis tasks, particularly for:
+- Pleural effusion status determination
+- Medical text similarity comparison
+- Clinical finding extraction
+
+### Sample Performance
+
+The model shows improved performance compared to the base model on medical text understanding tasks, particularly in distinguishing between different pleural effusion states and medical abbreviations.
+
+## Intended Use
 
+### Primary Use Cases
+- **Medical Text Embeddings**: Generate embeddings for chest X-ray reports
+- **Clinical Text Similarity**: Compare medical texts for semantic similarity
+- **Medical Information Retrieval**: Find relevant medical reports or findings
+- **Clinical NLP Research**: Foundation model for medical text analysis
+
+### Limitations
+- Specialized for chest X-ray reports - may not generalize to other medical domains
+- Requires careful preprocessing for optimal performance
+- Should be used as part of a larger clinical decision support system, not for standalone diagnosis
+
+## Technical Specifications
+
+- **Model Type**: Bidirectional Language Model (LLM2Vec)
+- **Architecture**: LlamaBiModel (modified Llama 3.2)
+- **Parameters**: ~1B parameters
+- **Input Length**: Up to 512 tokens
+- **Output**: Dense embeddings
+- **Precision**: bfloat16
+
+## Citation
+
+If you use this model in your research, please cite:
+
+```bibtex
+@misc{llm2vec4cxr,
+  title={LLM2Vec4CXR: Fine-tuned Language Model for Chest X-ray Report Analysis},
+  author={[Your Name]},
+  year={2024},
+  howpublished={\\url{https://huggingface.co/lukeingawesome/llm2vec4cxr}},
+}
 ```
-TBD
-```
+
+## Acknowledgments
+
+This model is built upon:
+- [LLM2Vec](https://github.com/McGill-NLP/llm2vec) - Framework for converting decoder-only LLMs into text encoders
+- [LLM2CLIP](https://github.com/microsoft/LLM2CLIP) - Microsoft's implementation for connecting LLMs with CLIP models
+- [microsoft/LLM2CLIP-Llama-3.2-1B-Instruct-CC-Finetuned](https://huggingface.co/microsoft/LLM2CLIP-Llama-3.2-1B-Instruct-CC-Finetuned) - Base model
+
+## License
+
+This model is licensed under the MIT License.

config.json

@@ -1,13 +1,14 @@
 {
-  "_name_or_path": "microsoft/LLM2CLIP-Llama-3.2-1B-Instruct-CC-Finetuned",
+  "_name_or_path": "lukeingawesome/llm2vec4cxr",
   "architectures": [
-    "LlamaBiForMNTP"
+    "LlamaBiModel"
   ],
   "attention_bias": false,
   "attention_dropout": 0.0,
   "auto_map": {
     "AutoModel": "microsoft/LLM2CLIP-Llama-3.2-1B-Instruct-CC-Finetuned--modeling_bidirectional_llama_encoder.LlamaBiModel"
   },
+  "base_model": "microsoft/LLM2CLIP-Llama-3.2-1B-Instruct-CC-Finetuned",
   "bos_token_id": 128000,
   "eos_token_id": [
     128001,
@@ -25,6 +26,7 @@
   "num_attention_heads": 32,
   "num_hidden_layers": 16,
   "num_key_value_heads": 8,
+  "pooling_mode": "latent_attention",
   "pretraining_tp": 1,
   "rms_norm_eps": 1e-05,
   "rope_scaling": {
@@ -39,5 +41,8 @@
   "torch_dtype": "bfloat16",
   "transformers_version": "4.44.2",
   "use_cache": true,
-  "vocab_size": 128256
-}
+  "vocab_size": 128256,
+  "model_description": "Fine-tuned LLM2Vec model for chest X-ray report analysis with latent attention pooling",
+  "domain": "medical",
+  "task": "text-embeddings"
+}

pytorch_model.bin

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

requirements.txt

@@ -0,0 +1,6 @@
+torch>=2.0.0
+transformers>=4.44.0
+accelerate>=0.20.0
+flash-attn>=2.5.0
+numpy>=1.21.0
+huggingface_hub>=0.16.0

special_tokens_map.json

@@ -13,6 +13,5 @@
     "rstrip": false,
     "single_word": false
   },
-  "mask_token": "_",
   "pad_token": "<|eot_id|>"
 }

tokenizer.json

@@ -1,17 +1,20 @@
 {
   "version": "1.0",
-  "truncation": null,
-  "padding": null,
+  "truncation": {
+    "direction": "Right",
+    "max_length": 512,
+    "strategy": "LongestFirst",
+    "stride": 0
+  },
+  "padding": {
+    "strategy": "BatchLongest",
+    "direction": "Left",
+    "pad_to_multiple_of": null,
+    "pad_id": 128009,
+    "pad_type_id": 0,
+    "pad_token": "<|eot_id|>"
+  },
   "added_tokens": [
-    {
-      "id": 62,
-      "content": "_",
-      "single_word": false,
-      "lstrip": false,
-      "rstrip": false,
-      "normalized": false,
-      "special": true
-    },
     {
       "id": 128000,
       "content": "<|begin_of_text|>",

tokenizer_config.json

@@ -1,13 +1,5 @@
 {
   "added_tokens_decoder": {
-    "62": {
-      "content": "_",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
     "128000": {
       "content": "<|begin_of_text|>",
       "lstrip": false,
@@ -2061,7 +2053,6 @@
   "chat_template": "{{- bos_token }}\n{%- if custom_tools is defined %}\n    {%- set tools = custom_tools %}\n{%- endif %}\n{%- if not tools_in_user_message is defined %}\n    {%- set tools_in_user_message = true %}\n{%- endif %}\n{%- if not date_string is defined %}\n    {%- if strftime_now is defined %}\n        {%- set date_string = strftime_now(\"%d %b %Y\") %}\n    {%- else %}\n        {%- set date_string = \"26 Jul 2024\" %}\n    {%- endif %}\n{%- endif %}\n{%- if not tools is defined %}\n    {%- set tools = none %}\n{%- endif %}\n\n{#- This block extracts the system message, so we can slot it into the right place. #}\n{%- if messages[0]['role'] == 'system' %}\n    {%- set system_message = messages[0]['content']|trim %}\n    {%- set messages = messages[1:] %}\n{%- else %}\n    {%- set system_message = \"\" %}\n{%- endif %}\n\n{#- System message #}\n{{- \"<|start_header_id|>system<|end_header_id|>\\n\\n\" }}\n{%- if tools is not none %}\n    {{- \"Environment: ipython\\n\" }}\n{%- endif %}\n{{- \"Cutting Knowledge Date: December 2023\\n\" }}\n{{- \"Today Date: \" + date_string + \"\\n\\n\" }}\n{%- if tools is not none and not tools_in_user_message %}\n    {{- \"You have access to the following functions. To call a function, please respond with JSON for a function call.\" }}\n    {{- 'Respond in the format {\"name\": function name, \"parameters\": dictionary of argument name and its value}.' }}\n    {{- \"Do not use variables.\\n\\n\" }}\n    {%- for t in tools %}\n        {{- t | tojson(indent=4) }}\n        {{- \"\\n\\n\" }}\n    {%- endfor %}\n{%- endif %}\n{{- system_message }}\n{{- \"<|eot_id|>\" }}\n\n{#- Custom tools are passed in a user message with some extra guidance #}\n{%- if tools_in_user_message and not tools is none %}\n    {#- Extract the first user message so we can plug it in here #}\n    {%- if messages | length != 0 %}\n        {%- set first_user_message = messages[0]['content']|trim %}\n        {%- set messages = messages[1:] %}\n    {%- else %}\n        {{- raise_exception(\"Cannot put tools in the first user message when there's no first user message!\") }}\n{%- endif %}\n    {{- '<|start_header_id|>user<|end_header_id|>\\n\\n' -}}\n    {{- \"Given the following functions, please respond with a JSON for a function call \" }}\n    {{- \"with its proper arguments that best answers the given prompt.\\n\\n\" }}\n    {{- 'Respond in the format {\"name\": function name, \"parameters\": dictionary of argument name and its value}.' }}\n    {{- \"Do not use variables.\\n\\n\" }}\n    {%- for t in tools %}\n        {{- t | tojson(indent=4) }}\n        {{- \"\\n\\n\" }}\n    {%- endfor %}\n    {{- first_user_message + \"<|eot_id|>\"}}\n{%- endif %}\n\n{%- for message in messages %}\n    {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}\n        {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\\n\\n'+ message['content'] | trim + '<|eot_id|>' }}\n    {%- elif 'tool_calls' in message %}\n        {%- if not message.tool_calls|length == 1 %}\n            {{- raise_exception(\"This model only supports single tool-calls at once!\") }}\n        {%- endif %}\n        {%- set tool_call = message.tool_calls[0].function %}\n        {{- '<|start_header_id|>assistant<|end_header_id|>\\n\\n' -}}\n        {{- '{\"name\": \"' + tool_call.name + '\", ' }}\n        {{- '\"parameters\": ' }}\n        {{- tool_call.arguments | tojson }}\n        {{- \"}\" }}\n        {{- \"<|eot_id|>\" }}\n    {%- elif message.role == \"tool\" or message.role == \"ipython\" %}\n        {{- \"<|start_header_id|>ipython<|end_header_id|>\\n\\n\" }}\n        {%- if message.content is mapping or message.content is iterable %}\n            {{- message.content | tojson }}\n        {%- else %}\n            {{- message.content }}\n        {%- endif %}\n        {{- \"<|eot_id|>\" }}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|start_header_id|>assistant<|end_header_id|>\\n\\n' }}\n{%- endif %}\n",
   "clean_up_tokenization_spaces": true,
   "eos_token": "<|eot_id|>",
-  "mask_token": "_",
   "model_input_names": [
     "input_ids",
     "attention_mask"

usage_example.py

@@ -0,0 +1,236 @@
+"""
+Example usage script for LLM2Vec4CXR model.
+This demonstrates how to load and use the model for chest X-ray report analysis.
+"""
+
+import torch
+import torch.nn.functional as F
+from llm2vec_wrapper import LLM2VecWrapper as LLM2Vec
+
+def load_llm2vec4cxr_model(model_name_or_path="lukeingawesome/llm2vec4cxr"):
+    """
+    Load the LLM2Vec4CXR model with proper configuration.
+    
+    Args:
+        model_name_or_path (str): Hugging Face model path or local path
+    
+    Returns:
+        tuple: (model, tokenizer)
+    """
+    # Load model with the specific configuration used for LLM2Vec4CXR
+    model = LLM2Vec.from_pretrained(
+        base_model_name_or_path=model_name_or_path,
+        enable_bidirectional=True,
+        pooling_mode="latent_attention",  # This is the key modification
+        max_length=512,
+        torch_dtype=torch.bfloat16,
+    )
+    
+    # Configure tokenizer
+    tokenizer = model.tokenizer
+    tokenizer.padding_side = 'left'
+    
+    return model, tokenizer
+
+def tokenize_with_separator(texts, tokenizer, max_length=512):
+    """
+    Tokenize texts with special handling for separator-based splitting.
+    This is useful for instruction-following tasks.
+    
+    Args:
+        texts (list): List of texts to tokenize
+        tokenizer: The tokenizer to use
+        max_length (int): Maximum sequence length
+    
+    Returns:
+        dict: Tokenized inputs with attention masks and embed masks
+    """
+    texts_2 = []
+    original_texts = []
+    separator = '!@#$%^&*()'
+    
+    for text in texts:
+        parts = text.split(separator)
+        texts_2.append(parts[1] if len(parts) > 1 else "")
+        original_texts.append("".join(parts))
+
+    # Tokenize original texts
+    tokenized = tokenizer(
+        original_texts,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=max_length,
+    )
+    
+    # Create embedding masks for the separated parts
+    embed_mask = None
+    for t_i, t in enumerate(texts_2):
+        ids = tokenizer(
+            [t],
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=max_length,
+            add_special_tokens=False,
+        )
+        
+        e_m = torch.zeros_like(tokenized["attention_mask"][t_i])
+        if len(ids["input_ids"][0]) > 0:
+            e_m[-len(ids["input_ids"][0]):] = torch.ones(len(ids["input_ids"][0]))
+            
+        if embed_mask is None:
+            embed_mask = e_m.unsqueeze(0)
+        else:
+            embed_mask = torch.cat((embed_mask, e_m.unsqueeze(0)), dim=0)
+
+    tokenized["embed_mask"] = embed_mask
+    return tokenized
+
+def compute_similarities(model, tokenizer, texts, device):
+    """
+    Compute similarity scores between the first text and all other texts.
+    
+    Args:
+        model: The LLM2Vec model
+        tokenizer: The tokenizer
+        texts (list): List of texts to compare (first text is the reference)
+        device: The device to run computations on
+    
+    Returns:
+        tuple: (embeddings, similarities)
+    """
+    with torch.no_grad():
+        # Use separator-based tokenization if texts contain the separator
+        if any('!@#$%^&*()' in text for text in texts):
+            tokenized = tokenize_with_separator(texts, tokenizer, 512)
+        else:
+            tokenized = tokenizer(
+                texts,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=512,
+            )
+        
+        tokenized = tokenized.to(device)
+        if hasattr(tokenized, 'to'):
+            tokenized = tokenized.to(torch.bfloat16)
+        else:
+            # Convert each tensor in the dict
+            for key in tokenized:
+                if torch.is_tensor(tokenized[key]):
+                    tokenized[key] = tokenized[key].to(torch.bfloat16)
+        
+        embeddings = model(tokenized)
+        
+        # Compute cosine similarities between first embedding and all others
+        similarities = F.cosine_similarity(embeddings[0], embeddings[1:], dim=1)
+    
+    return embeddings, similarities
+
+def main():
+    """
+    Example usage of the LLM2Vec4CXR model for chest X-ray report analysis.
+    """
+    # Set device
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Using device: {device}")
+    
+    # Load the model
+    print("Loading LLM2Vec4CXR model...")
+    model, tokenizer = load_llm2vec4cxr_model()
+    model = model.to(device).to(torch.bfloat16)
+    model.eval()
+    
+    # Example 1: Basic text embedding
+    print("\n" + "="*60)
+    print("Example 1: Basic Text Embedding")
+    print("="*60)
+    
+    report = "There is a small increase in the left-sided effusion. There continues to be volume loss at both bases."
+    
+    inputs = tokenizer(report, return_tensors="pt", padding=True, truncation=True, max_length=512)
+    inputs = inputs.to(device)
+    
+    with torch.no_grad():
+        embedding = model(inputs)
+    
+    print(f"Report: {report}")
+    print(f"Embedding shape: {embedding.shape}")
+    print(f"Embedding norm: {torch.norm(embedding).item():.4f}")
+    
+    # Example 2: Instruction-based similarity comparison
+    print("\n" + "="*60)
+    print("Example 2: Instruction-based Similarity Comparison")
+    print("="*60)
+    
+    separator = '!@#$%^&*()'
+    instruction = 'Determine the change or the status of the pleural effusion.'
+    report = 'There is a small increase in the left-sided effusion. There continues to be volume loss at both bases.'
+    text = instruction + separator + report
+    
+    comparison_options = [
+        'No pleural effusion',
+        'Pleural effusion',
+        'Effusion is seen in the right',
+        'Effusion is seen in the left',
+        'Pleural effusion is improving',
+        'Pleural effusion is stable',
+        'Pleural effusion is worsening'
+    ]
+    
+    all_texts = [text] + comparison_options
+    
+    # Compute similarities
+    _, similarities = compute_similarities(model, tokenizer, all_texts, device)
+    
+    print(f"Original text: {report}")
+    print(f"Instruction: {instruction}")
+    print("\nSimilarity Scores:")
+    print("-" * 50)
+    
+    for option, score in zip(comparison_options, similarities):
+        print(f"{option:<35} | {score.item():.4f}")
+    
+    # Find the most similar option
+    best_match_idx = torch.argmax(similarities).item()
+    print(f"\nBest match: {comparison_options[best_match_idx]} (score: {similarities[best_match_idx].item():.4f})")
+    
+    # Example 3: Multiple report comparison
+    print("\n" + "="*60)
+    print("Example 3: Multiple Report Comparison")
+    print("="*60)
+    
+    reports = [
+        "No acute cardiopulmonary abnormality.",
+        "Small bilateral pleural effusions.",
+        "Large left pleural effusion with compressive atelectasis.",
+        "Interval improvement in bilateral pleural effusions.",
+        "Worsening bilateral pleural effusions."
+    ]
+    
+    print("Computing embeddings for multiple reports...")
+    inputs = tokenizer(reports, return_tensors="pt", padding=True, truncation=True, max_length=512)
+    inputs = inputs.to(device)
+    
+    with torch.no_grad():
+        embeddings = model(inputs)
+    
+    # Compute pairwise similarities
+    similarity_matrix = F.cosine_similarity(
+        embeddings.unsqueeze(1), 
+        embeddings.unsqueeze(0), 
+        dim=2
+    )
+    
+    print("\nPairwise Similarity Matrix:")
+    print("-" * 30)
+    for i, report1 in enumerate(reports):
+        print(f"Report {i+1}: {report1[:30]}...")
+        for j, report2 in enumerate(reports):
+            print(f"  vs Report {j+1}: {similarity_matrix[i][j].item():.4f}")
+        print()
+
+if __name__ == "__main__":
+    main()