Add trust_remote_code integration (Qwen3-Embedding + LoRA)

.gitattributes

@@ -1,3 +1,4 @@
+<<<<<<< HEAD
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
@@ -33,3 +34,8 @@ 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
+=======
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+>>>>>>> 2dc7409 (Release chest2vec_0.6b_cxr (Stage2 LoRA + Stage3 pooler + API))

.gitignore

@@ -0,0 +1,24 @@
+# Build artifacts
+dist/
+build/
+*.egg-info/
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+.Python
+
+# Testing
+.pytest_cache/
+.coverage
+htmlcov/
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+
+# Jupyter
+.ipynb_checkpoints/
+

README.md

@@ -0,0 +1,215 @@
+---
+tags:
+- text-embeddings
+- retrieval
+- radiology
+- chest
+- qwen
+library_name: transformers
+---
+
+# chest2vec_0.6B
+
+This repository contains the *delta weights* for a global embedding model on top of **Qwen/Qwen3-Embedding-0.6B**:
+
+- **LoRA Adapter**: Contrastive LoRA adapter trained with multi-positive sigmoid loss under `./contrastive/`
+- **Inference helper**: `chest2vec.py`
+
+Base model weights are **not** included; they are downloaded from Hugging Face at runtime.
+
+## Model Architecture
+
+Chest2Vec is a two-stage model:
+1. **Base**: Qwen/Qwen3-Embedding-0.6B (downloaded at runtime)
+2. **LoRA Adapter**: Contrastive LoRA adapter trained with multi-positive sigmoid loss
+3. **Pooling**: Last-token pooling (EOS token) for global embeddings
+
+The model produces **global embeddings only** (no section-specific embeddings).
+
+## Installation
+
+Install the package and all dependencies:
+
+```bash
+# Install PyTorch with CUDA 12.6 support
+pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu126
+
+# Install transformers and trl
+pip install transformers==4.57.3 trl==0.9.3
+
+# Install deepspeed
+pip install deepspeed==0.16.9
+
+# Install flash-attention
+pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.6cxx11abiTRUE-cp310-cp310-linux_x86_64.whl
+
+# Install chest2vec package
+pip install chest2vec
+```
+
+Or use the installation script:
+
+```bash
+bash install_deps.sh
+```
+
+## Requirements
+
+This model **requires FlashAttention-2** (CUDA) by default, which is automatically installed with the package.
+
+## Quickstart
+
+### Installation + Loading
+
+```python
+from chest2vec import Chest2Vec
+
+# Load model from Hugging Face Hub
+m = Chest2Vec.from_pretrained("lukeingawesome/chest2vec_0.6b_chest", device="cuda:0")
+```
+
+### Instruction + Query Embeddings
+
+```python
+instructions = ["Find findings about the lungs."]
+queries = ["Consolidation in the right lower lobe."]
+
+out = m.embed_instruction_query(instructions, queries, max_len=512, batch_size=8)
+
+# Global embedding (last-token pooling)
+emb = out.embedding  # [N, H]
+```
+
+### Candidate Embeddings (Retrieval Bank)
+
+```python
+candidates = [
+    "Lungs are clear. No focal consolidation.",
+    "Pleural effusion on the left.",
+    "Cardiomediastinal silhouette is normal."
+]
+
+cand_out = m.embed_texts(candidates, max_len=512, batch_size=16)
+
+cand_emb = cand_out.embedding  # [N, H]
+```
+
+### Retrieval Example (Cosine Top-K)
+
+```python
+# Query embeddings
+q = out.embedding  # [Nq, H]
+
+# Document embeddings
+d = cand_out.embedding  # [Nd, H]
+
+# Compute top-k cosine similarities
+scores, idx = Chest2Vec.cosine_topk(q, d, k=5, device="cuda")
+# scores: [Nq, k] - similarity scores
+# idx: [Nq, k] - indices of top-k candidates
+
+print(f"Top-5 scores: {scores[0]}")
+print(f"Top-5 indices: {idx[0]}")
+```
+
+## API Reference
+
+### `Chest2Vec.from_pretrained()`
+
+Load the model from Hugging Face Hub or local path.
+
+```python
+m = Chest2Vec.from_pretrained(
+    repo_id_or_path: str,      # Hugging Face repo ID or local path
+    device: str = "cuda:0",    # Device to load model on
+    use_4bit: bool = False,    # Use 4-bit quantization
+    force_flash_attention_2: bool = True
+)
+```
+
+### `embed_instruction_query()`
+
+Embed instruction-query pairs. Returns `EmbedOutput` with:
+- `embedding`: `[N, H]` - global embeddings (L2-normalized, last-token pooling)
+
+```python
+out = m.embed_instruction_query(
+    instructions: List[str],
+    queries: List[str],
+    max_len: int = 512,
+    batch_size: int = 16
+)
+```
+
+### `embed_texts()`
+
+Embed plain texts (for document/candidate encoding).
+
+```python
+out = m.embed_texts(
+    texts: List[str],
+    max_len: int = 512,
+    batch_size: int = 16
+)
+```
+
+Returns `EmbedOutput` with:
+- `embedding`: `[N, H]` - global embeddings (L2-normalized, last-token pooling)
+
+### `cosine_topk()`
+
+Static method for efficient top-k cosine similarity search.
+
+```python
+scores, idx = Chest2Vec.cosine_topk(
+    query_emb: torch.Tensor,  # [Nq, H]
+    cand_emb: torch.Tensor,   # [Nd, H]
+    k: int = 10,
+    device: str = "cuda"
+)
+```
+
+## Model Files
+
+- `chest2vec.py` - Model class and inference utilities
+- `chest2vec_config.json` - Model configuration
+- `contrastive/` - LoRA adapter directory
+  - `adapter_config.json` - LoRA adapter configuration
+  - `adapter_model.safetensors` - LoRA adapter weights
+
+## Citation
+
+If you use this model, please cite:
+
+```bibtex
+@misc{chest2vec_0.6b_chest,
+  title={Chest2Vec: Global Embeddings for Chest X-Ray Reports},
+  author={Your Name},
+  year={2024},
+  howpublished={\url{https://huggingface.co/lukeingawesome/chest2vec_0.6b_chest}}
+}
+```
+
+## License
+
+[Specify your license here]
+
+## Usage (🤗 transformers)
+
+```python
+from transformers import AutoModel
+
+# base Qwen3-Embedding weights download automatically; needs trust_remote_code
+model = AutoModel.from_pretrained("chest2vec/chest2vec_0.6B", trust_remote_code=True)
+
+emb = model.embed_texts([
+    "Frontal chest radiograph. No focal consolidation. No pneumothorax. Heart size normal.",
+])
+emb            # [N, H] L2-normalized report embedding (last-token / EOS pooling)
+
+# similarity
+(emb[0] @ emb[1])    # cosine similarity (rows are unit-norm)
+```
+
+FlashAttention-2 is used automatically on CUDA when `flash-attn>=2` is installed
+(matching training); otherwise it falls back to SDPA so the model also loads on CPU.

config.json

@@ -0,0 +1,14 @@
+{
+  "model_type": "chest2vec",
+  "architectures": [
+    "Chest2VecModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_chest2vec.Chest2VecConfig",
+    "AutoModel": "modeling_chest2vec.Chest2VecModel"
+  },
+  "base_model": "Qwen/Qwen3-Embedding-0.6B",
+  "adapter_subdir": "contrastive",
+  "require_flash_attention_2": true,
+  "default_max_len": 512
+}

configuration_chest2vec.py

@@ -0,0 +1,26 @@
+"""Configuration for Chest2Vec — a LoRA-tuned Qwen3-Embedding model for
+chest radiology report embeddings.
+
+Chest2Vec = Qwen3-Embedding base + contrastive LoRA adapter. It produces a
+single L2-normalized report embedding (last-token / EOS pooling), matching the
+Qwen3-Embedding convention.
+"""
+from transformers import PretrainedConfig
+
+
+class Chest2VecConfig(PretrainedConfig):
+    model_type = "chest2vec"
+
+    def __init__(
+        self,
+        base_model: str = "Qwen/Qwen3-Embedding-0.6B",
+        adapter_subdir: str = "contrastive",
+        require_flash_attention_2: bool = True,
+        default_max_len: int = 512,
+        **kwargs,
+    ):
+        self.base_model = base_model
+        self.adapter_subdir = adapter_subdir
+        self.require_flash_attention_2 = require_flash_attention_2
+        self.default_max_len = default_max_len
+        super().__init__(**kwargs)

contrastive/adapter_config.json

@@ -0,0 +1,42 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": {
+    "base_model_class": "Qwen3Model",
+    "parent_library": "transformers.models.qwen3.modeling_qwen3"
+  },
+  "base_model_name_or_path": "Qwen/Qwen3-Embedding-0.6B",
+  "bias": "none",
+  "corda_config": null,
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 32,
+  "lora_bias": false,
+  "lora_dropout": 0.1,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 16,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "up_proj",
+    "k_proj",
+    "o_proj",
+    "v_proj",
+    "gate_proj",
+    "q_proj",
+    "down_proj"
+  ],
+  "task_type": null,
+  "trainable_token_indices": null,
+  "use_dora": false,
+  "use_rslora": false
+}

contrastive/adapter_model.safetensors

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

modeling_chest2vec.py

@@ -0,0 +1,301 @@
+"""Chest2Vec — LoRA-tuned Qwen3-Embedding model for chest radiology reports.
+
+Load with:
+
+    from transformers import AutoModel
+    model = AutoModel.from_pretrained("chest2vec/chest2vec_0.6B", trust_remote_code=True)
+    emb = model.embed_texts(["Frontal chest radiograph. No pneumothorax."])  # [N, H], L2-normalized
+
+Architecture:
+  1. Base   : Qwen/Qwen3-Embedding-{0.6B,4B}  (downloaded at runtime)
+  2. Adapter: frozen contrastive LoRA adapter  (./contrastive)
+
+Embeddings use last-token (EOS) pooling with left padding, matching Qwen3-Embedding
+and the Stage-2 training setup. FlashAttention-2 is used when CUDA + flash-attn>=2
+are available (matching training); otherwise it falls back to SDPA so the model
+also loads on CPU.
+"""
+import os
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn.functional as F
+
+from transformers import AutoTokenizer, AutoModel, BitsAndBytesConfig, PreTrainedModel
+
+from .configuration_chest2vec import Chest2VecConfig
+
+try:
+    from peft import PeftModel
+    _HAS_PEFT = True
+except Exception:
+    PeftModel = None
+    _HAS_PEFT = False
+
+try:
+    from huggingface_hub import snapshot_download
+    _HAS_HUB = True
+except Exception:
+    snapshot_download = None
+    _HAS_HUB = False
+
+
+# ----------------------------------------------------------------------------
+# Attention backend selection
+# ----------------------------------------------------------------------------
+def _flash_attn_available() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    try:
+        import flash_attn  # noqa: F401
+        ver = getattr(flash_attn, "__version__", "0.0.0")
+        return int(str(ver).split(".")[0]) >= 2
+    except Exception:
+        return False
+
+
+def _pick_attn_impl(requested: Optional[str], want_flash: bool) -> str:
+    import warnings
+    if requested:
+        return requested
+    if want_flash and _flash_attn_available():
+        return "flash_attention_2"
+    if want_flash:
+        warnings.warn(
+            "Chest2Vec was trained with FlashAttention-2, but it is unavailable "
+            "(needs CUDA + flash-attn>=2). Falling back to 'sdpa'; embeddings may "
+            "differ very slightly from the reference implementation.",
+            RuntimeWarning,
+        )
+    return "sdpa"
+
+
+# ----------------------------------------------------------------------------
+# Tokenization / pooling helpers (match Qwen3-Embedding + training)
+# ----------------------------------------------------------------------------
+def build_qwen_query(instruction: str, query: str) -> str:
+    return f"Instruct: {str(instruction).strip()}\nQuery: {str(query).strip()}"
+
+
+def get_pool_token_id(tok) -> int:
+    eod_id = tok.convert_tokens_to_ids("<|endoftext|>")
+    if eod_id is None or eod_id < 0:
+        eod_id = tok.pad_token_id
+    return eod_id
+
+
+def encode_with_eos_ids(tok, texts: List[str], max_len: int) -> Dict[str, torch.Tensor]:
+    """add_special_tokens=False, truncate to max_len-1, append <|endoftext|>, left-pad."""
+    pad_id = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
+    eod_id = get_pool_token_id(tok)
+    enc = tok(
+        [str(t) for t in texts],
+        add_special_tokens=False,
+        truncation=True,
+        max_length=max_len - 1,
+        padding=False,
+        return_attention_mask=False,
+    )
+    input_ids = [ids + [eod_id] for ids in enc["input_ids"]]
+    attn_mask = [[1] * len(ids) for ids in input_ids]
+    T = max((len(ids) for ids in input_ids), default=1)
+    input_ids = [[pad_id] * (T - len(ids)) + ids for ids in input_ids]
+    attn_mask = [[0] * (T - len(m)) + m for m in attn_mask]
+    return {
+        "input_ids": torch.tensor(input_ids, dtype=torch.long),
+        "attention_mask": torch.tensor(attn_mask, dtype=torch.long),
+    }
+
+
+def last_token_pool(last_hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+    """Left-padding-aware last-token (EOS) pooling."""
+    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+    if left_padding:
+        return last_hidden_states[:, -1]
+    idx = attention_mask.sum(dim=1) - 1
+    return last_hidden_states[torch.arange(last_hidden_states.size(0), device=last_hidden_states.device), idx]
+
+
+def get_last_hidden_state(model, input_ids, attention_mask):
+    m = model.module if hasattr(model, "module") else model
+    position_ids = attention_mask.long().cumsum(-1) - 1
+    position_ids.masked_fill_(attention_mask == 0, 0)
+    out = m(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids,
+            use_cache=False, return_dict=True)
+    if getattr(out, "last_hidden_state", None) is not None:
+        return out.last_hidden_state
+    out = m(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids,
+            output_hidden_states=True, use_cache=False, return_dict=True)
+    return out.hidden_states[-1]
+
+
+class Chest2VecModel(PreTrainedModel):
+    """LoRA-tuned Qwen3-Embedding model producing L2-normalized report embeddings."""
+
+    config_class = Chest2VecConfig
+    base_model_prefix = "chest2vec"
+    # Attention is handled by the inner Qwen3 backbone; advertise support so the
+    # transformers attn-implementation validator on this wrapper passes.
+    _supports_sdpa = True
+    _supports_flash_attn_2 = True
+    _supports_flash_attn = True
+    _supports_attention_backend = True
+
+    def __init__(self, config: Chest2VecConfig):
+        super().__init__(config)
+        # The base+adapter are assembled in `from_pretrained` (base downloads at runtime).
+        self.backbone = None
+        self.tokenizer = None
+        self._device = torch.device("cpu")
+        self.register_buffer("_anchor", torch.zeros(1), persistent=False)
+
+    def get_input_embeddings(self):
+        return None
+
+    def set_input_embeddings(self, value):
+        pass
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        config = kwargs.pop("config", None)
+        device = kwargs.pop("device", None)
+        use_4bit = kwargs.pop("use_4bit", False)
+        attn_implementation = kwargs.pop("attn_implementation", None)
+        torch_dtype = kwargs.pop("torch_dtype", None)
+        token = kwargs.pop("token", None) or kwargs.pop("use_auth_token", None)
+        cache_dir = kwargs.pop("cache_dir", None)
+        # remaining HF plumbing kwargs (state_dict, low_cpu_mem_usage, ...) are ignored
+
+        repo_path = pretrained_model_name_or_path
+        if not os.path.isdir(repo_path):
+            if not _HAS_HUB:
+                raise RuntimeError("huggingface_hub is required to load by repo_id.")
+            repo_path = snapshot_download(repo_path, token=token, cache_dir=cache_dir)
+
+        if config is None:
+            config = Chest2VecConfig.from_pretrained(repo_path)
+
+        if device is None:
+            device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        device_t = torch.device(device)
+        if torch_dtype is None:
+            torch_dtype = torch.bfloat16 if device_t.type == "cuda" else torch.float32
+
+        model = cls(config)
+        model._assemble(repo_path, device=device_t, use_4bit=use_4bit,
+                        attn_implementation=attn_implementation, torch_dtype=torch_dtype, token=token)
+        return model
+
+    def _assemble(self, repo_path, *, device, use_4bit, attn_implementation, torch_dtype, token=None):
+        cfg = self.config
+        if not _HAS_PEFT:
+            raise RuntimeError("peft is required. Install: pip install peft")
+
+        attn_impl = _pick_attn_impl(attn_implementation, bool(cfg.require_flash_attention_2))
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            cfg.base_model, padding_side="left", trust_remote_code=True, token=token
+        )
+        if tokenizer.pad_token_id is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        base_kwargs = dict(trust_remote_code=True, attn_implementation=attn_impl, token=token)
+        if use_4bit:
+            base_kwargs["quantization_config"] = BitsAndBytesConfig(
+                load_in_4bit=True, bnb_4bit_quant_type="nf4",
+                bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16,
+            )
+            base_kwargs["device_map"] = {"": str(device)}
+        else:
+            base_kwargs["torch_dtype"] = torch_dtype
+            if device.type == "cuda":
+                base_kwargs["device_map"] = {"": str(device)}
+        try:
+            base = AutoModel.from_pretrained(cfg.base_model, **base_kwargs)
+        except TypeError as e:
+            raise RuntimeError("transformers too old for attn_implementation=...; please upgrade.") from e
+        if device.type != "cuda" and not use_4bit:
+            base = base.to(device)
+
+        adapter_dir = os.path.join(repo_path, cfg.adapter_subdir)
+        if not os.path.isfile(os.path.join(adapter_dir, "adapter_config.json")):
+            raise FileNotFoundError(f"adapter_config.json not found under: {adapter_dir}")
+        backbone = PeftModel.from_pretrained(base, adapter_dir)
+        backbone.eval()
+
+        self.backbone = backbone
+        self.tokenizer = tokenizer
+        self._device = device
+        self.eval()
+
+    @property
+    def device(self):
+        return self._device
+
+    @torch.inference_mode()
+    def embed_texts(self, texts: List[str], *, max_len: Optional[int] = None,
+                    batch_size: int = 16, return_cpu_float32: bool = True) -> torch.Tensor:
+        """Return L2-normalized report embeddings, shape [N, H]."""
+        if self.backbone is None:
+            raise RuntimeError("Model not assembled; load via from_pretrained(...).")
+        max_len = int(max_len or self.config.default_max_len)
+        device = self._device
+        if device.type == "cuda":
+            amp_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+            use_amp = True
+        else:
+            amp_dtype, use_amp = torch.float32, False
+
+        outs = []
+        for i in range(0, len(texts), batch_size):
+            chunk = [str(t) for t in texts[i:i + batch_size]]
+            enc = encode_with_eos_ids(self.tokenizer, chunk, max_len)
+            input_ids = enc["input_ids"].to(device, non_blocking=True)
+            attention_mask = enc["attention_mask"].to(device, non_blocking=True)
+            with torch.autocast(device_type=("cuda" if device.type == "cuda" else "cpu"),
+                                dtype=amp_dtype, enabled=use_amp):
+                h = get_last_hidden_state(self.backbone, input_ids, attention_mask)
+                emb = F.normalize(last_token_pool(h, attention_mask).float(), p=2, dim=-1)
+            outs.append(emb.detach())
+        embeddings = torch.cat(outs, dim=0)
+        if return_cpu_float32:
+            embeddings = F.normalize(embeddings.float().cpu(), p=2, dim=-1)
+        return embeddings
+
+    @torch.inference_mode()
+    def embed_instruction_query(self, instructions: List[str], queries: List[str], **kw) -> torch.Tensor:
+        if len(instructions) != len(queries):
+            raise ValueError("instructions and queries must have the same length.")
+        return self.embed_texts([build_qwen_query(i, q) for i, q in zip(instructions, queries)], **kw)
+
+    def forward(self, texts: List[str], **kw) -> torch.Tensor:  # type: ignore[override]
+        return self.embed_texts(texts, **kw)
+
+    @staticmethod
+    def cosine_topk(query_emb, cand_emb, k=10, *, device="cuda",
+                    query_batch_size=256, doc_chunk_size=8192):
+        device_t = torch.device(device if torch.cuda.is_available() else "cpu")
+        q = F.normalize(query_emb.float(), p=2, dim=-1)
+        d = F.normalize(cand_emb.float(), p=2, dim=-1)
+        Nq, _ = q.shape
+        Nd = d.shape[0]
+        k = min(int(k), Nd)
+        top_scores_all = torch.empty((Nq, k), dtype=torch.float32)
+        top_indices_all = torch.empty((Nq, k), dtype=torch.long)
+        for qs in range(0, Nq, query_batch_size):
+            qe = q[qs:qs + query_batch_size].to(device_t, non_blocking=True)
+            bq = qe.size(0)
+            top_scores = torch.full((bq, k), -1e9, device=device_t, dtype=torch.float32)
+            top_indices = torch.full((bq, k), -1, device=device_t, dtype=torch.long)
+            for ds in range(0, Nd, doc_chunk_size):
+                de = d[ds:ds + doc_chunk_size].to(device_t, non_blocking=True)
+                scores = (qe @ de.T).float()
+                chunk = scores.size(1)
+                idx_chunk = torch.arange(ds, ds + chunk, device=device_t, dtype=torch.long).unsqueeze(0).expand(bq, -1)
+                comb_scores = torch.cat([top_scores, scores], dim=1)
+                comb_idx = torch.cat([top_indices, idx_chunk], dim=1)
+                new_scores, new_pos = torch.topk(comb_scores, k, dim=1)
+                top_scores, top_indices = new_scores, comb_idx.gather(1, new_pos)
+            top_scores_all[qs:qs + bq] = top_scores.cpu()
+            top_indices_all[qs:qs + bq] = top_indices.cpu()
+        return top_scores_all, top_indices_all

requirements.txt

@@ -0,0 +1,12 @@
+torch==2.6.0
+torchvision==0.21.0
+torchaudio==2.6.0
+transformers==4.57.3
+trl==0.9.3
+deepspeed==0.16.9
+https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.6cxx11abiTRUE-cp310-cp310-linux_x86_64.whl
+peft
+huggingface_hub
+bitsandbytes
+accelerate
+numpy