From 959561791f824a2591a3988d2e5a629cc7a9b605 Mon Sep 17 00:00:00 2001
From: jingyaogong <gongjy.cs@foxmail.com>
Date: Sun, 28 Jun 2026 18:13:08 +0800
Subject: [PATCH] [fix] robustness

---
 .gitignore              |  5 ++++-
 dataset/lm_dataset.py   | 10 ++++++----
 eval_vlm.py             |  4 +++-
 scripts/web_demo_vlm.py |  9 ++++++---
 4 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/.gitignore b/.gitignore
index 69225a3..e001a86 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,8 @@
 model/__pycache__
+__pycache__/
+*.py[cod]
+.DS_Store
 out
 website/
 docs-minimind-v/
-测试.py
\ No newline at end of file
+测试.py
diff --git a/dataset/lm_dataset.py b/dataset/lm_dataset.py
index a22a1d5..f8db4b0 100644
--- a/dataset/lm_dataset.py
+++ b/dataset/lm_dataset.py
@@ -9,6 +9,7 @@ import io
 from PIL import Image
 from torch.utils.data import Dataset, DataLoader
 from model.model_vlm import MiniMindVLM
+from datasets import Dataset as HFDataset
 import pyarrow as pa
 import pyarrow.parquet as pq
 
@@ -47,7 +48,7 @@ def post_processing_chat(prompt_content, empty_think_ratio=0.2):
 class VLMDataset(Dataset):
     def __init__(self, parquet_path, tokenizer, preprocess=None, max_length=512, image_special_token='<|image_pad|>', image_token_len=64):
         super().__init__()
-        self.table = pa.Table.from_batches(pq.ParquetFile(parquet_path).iter_batches())
+        self.dataset = HFDataset.from_parquet(parquet_path)
         self.tokenizer = tokenizer
         self.max_length = max_length
         self.preprocess = preprocess
@@ -56,7 +57,7 @@ class VLMDataset(Dataset):
         self.eos_id = tokenizer(f'{tokenizer.eos_token}\n', add_special_tokens=False).input_ids
 
     def __len__(self):
-        return len(self.table)
+        return len(self.dataset)
 
     def create_chat_prompt(self, conversations):
         messages = []
@@ -90,8 +91,9 @@ class VLMDataset(Dataset):
         return labels
 
     def __getitem__(self, index: int):
-        conversations = json.loads(self.table['conversations'][index].as_py())
-        image_bytes = self.table['image_bytes'][index].as_py()
+        row = self.dataset[index]
+        conversations = json.loads(row['conversations']) if isinstance(row['conversations'], str) else row['conversations']
+        image_bytes = row['image_bytes']
         if not isinstance(image_bytes, list): image_bytes = [image_bytes]
         
         conversations = pre_processing_chat(conversations)
diff --git a/eval_vlm.py b/eval_vlm.py
index 30b95b0..27accf4 100644
--- a/eval_vlm.py
+++ b/eval_vlm.py
@@ -25,7 +25,9 @@ def init_model(args):
         model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True)
         model.vision_encoder, model.processor = MiniMindVLM.get_vision_model("./model/siglip2-base-p32-256-ve")
     get_model_params(model, model.config)
-    return model.half().eval().to(args.device), tokenizer, model.processor
+    model = model.eval()
+    if "cuda" in args.device: model = model.half()
+    return model.to(args.device), tokenizer, model.processor
 
 
 def main():
diff --git a/scripts/web_demo_vlm.py b/scripts/web_demo_vlm.py
index cb54d71..19a466c 100644
--- a/scripts/web_demo_vlm.py
+++ b/scripts/web_demo_vlm.py
@@ -43,7 +43,9 @@ def load_vlm_model(model_path):
         if model.vision_encoder is None: raise FileNotFoundError(f"视觉编码器未找到: {vision_model_path}")
         preprocess = model.processor
         lm_config = model.config
-        model = model.half().eval().to(device)
+        model = model.eval()
+        if "cuda" in device: model = model.half()
+        model = model.to(device)
         model.vision_encoder = model.vision_encoder.to(device)
         current_model_name = os.path.basename(model_path)
         param_str = f'{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.2f}M'
@@ -76,13 +78,14 @@ def chat(prompt, current_image_path=None):
         messages,
         tokenize=False,
         add_generation_prompt=True
-    )[-max_seq_len + 1:]
+    )
 
     with torch.no_grad():
         inputs = tokenizer(
             new_prompt,
             return_tensors="pt",
-            truncation=True
+            truncation=True,
+            max_length=max_seq_len
         ).to(device)
         queue = Queue()
         streamer = CustomStreamer(tokenizer, queue)