added phi

app.py

@@ -215,6 +215,8 @@ def http_bot(
             template_name = "mpt_text"
         elif "llama-2" in model_name:
             template_name = "llama_2"
+        elif "phi" in model_name:
+            template_name = "phi"
         else:
             template_name = "vicuna_v1"
         new_state = conv_templates[template_name].copy()
@@ -604,7 +606,7 @@ if __name__ == "__main__":
     args = get_args()
     logger.info(f"args: {args}")
 
-    model_path = "liuhaotian/llava-v1.5-13b"
+    model_path = "marianna13/llava-phi-2-3b"
     bits = int(os.getenv("bits", 8))
 
     controller_proc = start_controller()

llava/__init__.py

@@ -1 +1,2 @@
 from .model import LlavaLlamaForCausalLM
+from .model import LlavaMistralForCausalLM

llava/conversation.py

@@ -10,6 +10,7 @@ class SeparatorStyle(Enum):
     MPT = auto()
     PLAIN = auto()
     LLAMA_2 = auto()
+    PHI = auto()
 
 
 @dataclasses.dataclass
@@ -72,6 +73,28 @@ class Conversation:
             wrap_inst = lambda msg: f"[INST] {msg} [/INST]"
             ret = ""
 
+            for i, (role, message) in enumerate(messages):
+                if i == 0:
+                    assert message, "first message should not be none"
+                    assert role == self.roles[0], "first message should come from user"
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    if i == 0: message = wrap_sys(self.system) + message
+                    if i % 2 == 0:
+                        message = wrap_inst(message)
+                        ret += self.sep + message
+                    else:
+                        ret += " " + message + " " + self.sep2
+                else:
+                    ret += ""
+            ret = ret.lstrip(self.sep)
+
+        elif self.sep_style == SeparatorStyle.PHI:
+            wrap_sys = lambda msg: f"<<SYS>>\n{msg}\n<</SYS>>\n\n"
+            wrap_inst = lambda msg: f"Instruct: {msg} \nOutput:"
+            ret = ""
+
             for i, (role, message) in enumerate(messages):
                 if i == 0:
                     assert message, "first message should not be none"
@@ -261,6 +284,30 @@ conv_vicuna_v1 = Conversation(
     sep2="</s>",
 )
 
+phi = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v0",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.PHI,
+    sep="<|endoftext|>",
+    sep2="<|endoftext|>",
+)
+
+conv_phi = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep="<|endoftext|>",
+    sep2="<|endoftext|>",
+)
+
 conv_llama_2 = Conversation(
     system="""You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
 
@@ -287,6 +334,19 @@ conv_llava_llama_2 = Conversation(
     sep2="</s>",
 )
 
+llava_phi = Conversation(
+    system="You are a helpful language and vision assistant. "
+           "You are able to understand the visual content that the user provides, "
+           "and assist the user with a variety of tasks using natural language.",
+    roles=("USER", "ASSISTANT"),
+    version="llava_phi",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep="<|endoftext|>",
+    sep2="<|endoftext|>",
+)
+
 conv_mpt = Conversation(
     system="""<|im_start|>system
 A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.""",
@@ -344,6 +404,18 @@ conv_llava_v1 = Conversation(
     sep2="</s>",
 )
 
+conv_mistral_v1 = Conversation(
+    system="""You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
+    If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.""",
+    roles=("user", "assistant"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
 conv_llava_v1_mmtag = Conversation(
     system="A chat between a curious user and an artificial intelligence assistant. "
            "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
@@ -364,7 +436,7 @@ conv_templates = {
     "v1": conv_vicuna_v1,
     "vicuna_v1": conv_vicuna_v1,
     "llama_2": conv_llama_2,
-
+    "mistral": conv_llama_2,
     "plain": conv_llava_plain,
     "v0_plain": conv_llava_plain,
     "llava_v0": conv_llava_v0,
@@ -372,7 +444,10 @@ conv_templates = {
     "llava_v1": conv_llava_v1,
     "v1_mmtag": conv_llava_v1_mmtag,
     "llava_llama_2": conv_llava_llama_2,
-
+    "conv_mistral_v1": conv_mistral_v1,
+    "llava_phi": llava_phi,
+    "conv_phi": conv_phi,
+    "phi": phi,
     "mpt": conv_mpt,
 }
 

llava/eval/eval_science_qa.py

@@ -13,6 +13,8 @@ def get_args():
     parser.add_argument('--output-result', type=str)
     parser.add_argument('--split', type=str, default='test')
     parser.add_argument('--options', type=list, default=["A", "B", "C", "D", "E"])
+    parser.add_argument('--pid-splits-path')
+    parser.add_argument('--problems-path')
     return parser.parse_args()
 
 
@@ -39,8 +41,8 @@ if __name__ == "__main__":
     args = get_args()
 
     base_dir = args.base_dir
-    split_indices = json.load(open(os.path.join(base_dir, "pid_splits.json")))[args.split]
-    problems = json.load(open(os.path.join(base_dir, "problems.json")))
+    split_indices = json.load(open(args.pid_splits_path))[args.split]
+    problems = json.load(open(args.problems_path))
     predictions = [json.loads(line) for line in open(args.result_file)]
     predictions = {pred['question_id']: pred for pred in predictions}
     split_problems = {idx: problems[idx] for idx in split_indices}
@@ -54,18 +56,26 @@ if __name__ == "__main__":
     sqa_results['outputs'] = {}
 
     for prob_id, prob in split_problems.items():
+        # prob_id = f'{args.split}/{prob_id}'
         if prob_id not in predictions:
-            continue
-        pred = predictions[prob_id]
-        pred_text = pred['text']
-
-        pattern = re.compile(r'The answer is ([A-Z]).')
-        res = pattern.findall(pred_text)
-        if len(res) == 1:
-            answer = res[0]  # 'A', 'B', ...
+            pred = {'text': 'FAILED', 'prompt': 'Unknown'}
+            pred_text = 'FAILED'
         else:
-            answer = "FAILED"
+            pred = predictions[prob_id]
+            pred_text = pred['text']
 
+        if pred_text in args.options:
+            answer = pred_text
+        elif len(pred_text) >= 3 and pred_text[0] in args.options and pred_text[1:3] == ". ":
+            answer = pred_text[0]
+        else:
+            pattern = re.compile(r'The answer is ([A-Z])')
+            res = pattern.findall(pred_text)
+            if len(res) == 1:
+                answer = res[0]  # 'A', 'B', ...
+            else:
+                answer = "FAILED"
+    
         pred_idx = get_pred_idx(answer, prob['choices'], args.options)
 
         analysis = {
@@ -85,9 +95,16 @@ if __name__ == "__main__":
         else:
             results['incorrect'].append(analysis)
 
+
     correct = len(results['correct'])
     total = len(results['correct']) + len(results['incorrect'])
-    print(f'Total: {total}, Correct: {correct}, Accuracy: {correct / total * 100:.2f}%')
+    ###### IMG ######
+    multimodal_correct = len([x for x in results['correct'] if x['is_multimodal']])
+    multimodal_incorrect = len([x for x in results['incorrect'] if x['is_multimodal']])
+    multimodal_total = multimodal_correct + multimodal_incorrect
+    ###### IMG ######
+
+    print(f'Total: {total}, Correct: {correct}, Accuracy: {correct / total * 100:.2f}%, IMG-Accuracy: {multimodal_correct / multimodal_total * 100:.2f}%')
 
     sqa_results['acc'] = correct / total * 100
     sqa_results['correct'] = correct
@@ -96,4 +113,4 @@ if __name__ == "__main__":
     with open(args.output_file, 'w') as f:
         json.dump(results, f, indent=2)
     with open(args.output_result, 'w') as f:
-        json.dump(sqa_results, f, indent=2)
+        json.dump(sqa_results, f, indent=2)

llava/eval/model_vqa_science.py

@@ -57,6 +57,10 @@ def eval_model(args):
         else:
             images = None
 
+        if args.single_pred_prompt:
+            qs = qs + '\n' + "Answer with the option's letter from the given choices directly."
+            cur_prompt = cur_prompt + '\n' + "Answer with the option's letter from the given choices directly."
+
         conv = conv_templates[args.conv_mode].copy()
         conv.append_message(conv.roles[0], qs)
         conv.append_message(conv.roles[1], None)
@@ -64,19 +68,22 @@ def eval_model(args):
 
         input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
 
-        stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
+        stop_str = conv.sep2
+        stop_str = "\n" if "phi" in model_name else stop_str
         keywords = [stop_str]
-        stopping_criteria = [KeywordsStoppingCriteria(keywords, tokenizer, input_ids)] if conv.version == "v0" else None
+        stopping_criteria = [KeywordsStoppingCriteria(keywords, tokenizer, input_ids)]
+        eos_token_id = tokenizer.eos_token_id
 
         with torch.inference_mode():
             output_ids = model.generate(
                 input_ids,
                 images=images,
-                do_sample=True,
-                temperature=0.2,
+                do_sample=True if args.temperature > 0 else False,
+                temperature=args.temperature,
                 max_new_tokens=1024,
                 use_cache=True,
                 stopping_criteria=stopping_criteria,
+                # eos_token_id=eos_token_id
             )
 
         input_token_len = input_ids.shape[1]
@@ -88,7 +95,9 @@ def eval_model(args):
         if outputs.endswith(stop_str):
             outputs = outputs[:-len(stop_str)]
         outputs = outputs.strip()
-
+        # outputs = outputs.replace("\n<</SYS>>", "")
+        # print("question:\n", cur_prompt)
+        # print("answer:\n", outputs)
         # prompt for answer
         if args.answer_prompter:
             outputs_reasoning = outputs
@@ -98,11 +107,11 @@ def eval_model(args):
                 output_ids = model.generate(
                     input_ids,
                     images=images,
-                    do_sample=True,
-                    temperature=0.2,
+                    do_sample=True if args.temperature > 0 else False,
+                    temperature=args.temperature,
                     max_new_tokens=64,
                     use_cache=True,
-                    stopping_criteria=[stopping_criteria])
+                    stopping_criteria=stopping_criteria)
 
             input_token_len = input_ids.shape[1]
             n_diff_input_output = (input_ids != output_ids[:, :input_token_len]).sum().item()
@@ -135,7 +144,9 @@ if __name__ == "__main__":
     parser.add_argument("--conv-mode", type=str, default="llava_v0")
     parser.add_argument("--num-chunks", type=int, default=1)
     parser.add_argument("--chunk-idx", type=int, default=0)
+    parser.add_argument("--temperature", type=float, default=0.2)
     parser.add_argument("--answer-prompter", action="store_true")
+    parser.add_argument("--single-pred-prompt", action="store_true")
     args = parser.parse_args()
 
-    eval_model(args)
+    eval_model(args)

llava/model/__init__.py

@@ -1,2 +1,4 @@
 from .language_model.llava_llama import LlavaLlamaForCausalLM, LlavaConfig
-from .language_model.llava_mpt import LlavaMPTForCausalLM, LlavaMPTConfig
+from .language_model.llava_mistral import LlavaMistralForCausalLM, LlavaConfig
+from .language_model.llava_phi import LlavaPhiForCausalLM, LlavaConfig
+# from .language_model.llava_mpt import LlavaMPTForCausalLM, LlavaMPTConfig

llava/model/apply_delta.py

@@ -8,6 +8,7 @@ import torch
 from tqdm import tqdm
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from llava import LlavaLlamaForCausalLM
+from llava import LlavaMistralForCausalLM
 
 
 def apply_delta(base_model_path, target_model_path, delta_path):
@@ -16,7 +17,7 @@ def apply_delta(base_model_path, target_model_path, delta_path):
         base_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
 
     print("Loading delta")
-    delta = LlavaLlamaForCausalLM.from_pretrained(delta_path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
+    delta = LlavaMistralForCausalLM.from_pretrained(delta_path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
     delta_tokenizer = AutoTokenizer.from_pretrained(delta_path)
 
     print("Applying delta")

llava/model/builder.py

@@ -25,7 +25,6 @@ from llava.constants import DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, D
 
 def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, load_4bit=False, device_map="auto"):
     kwargs = {"device_map": device_map}
-    kwargs["offload_folder"] = "offload"
 
     if load_8bit:
         kwargs['load_in_8bit'] = True
@@ -48,7 +47,7 @@ def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, l
             lora_cfg_pretrained = AutoConfig.from_pretrained(model_path)
             tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
             print('Loading LLaVA from base model...')
-            model = LlavaLlamaForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=lora_cfg_pretrained, **kwargs)
+            model = LlavaMistralForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=lora_cfg_pretrained, **kwargs)
             token_num, tokem_dim = model.lm_head.out_features, model.lm_head.in_features
             if model.lm_head.weight.shape[0] != token_num:
                 model.lm_head.weight = torch.nn.Parameter(torch.empty(token_num, tokem_dim, device=model.device, dtype=model.dtype))
@@ -90,18 +89,24 @@ def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, l
             else:
                 tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
                 cfg_pretrained = AutoConfig.from_pretrained(model_path)
-                model = LlavaLlamaForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=cfg_pretrained, **kwargs)
+                model = LlavaMistralForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=cfg_pretrained, **kwargs)
 
             mm_projector_weights = torch.load(os.path.join(model_path, 'mm_projector.bin'), map_location='cpu')
             mm_projector_weights = {k: v.to(torch.float16) for k, v in mm_projector_weights.items()}
             model.load_state_dict(mm_projector_weights, strict=False)
         else:
-            if 'mpt' in model_name.lower():
+            if 'phi' in model_name.lower():
+                tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+                model = LlavaPhiForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
+            elif 'mpt' in model_name.lower():
                 tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
                 model = LlavaMPTForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
+            elif 'phi' in model_name.lower():
+                tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+                model = LlavaPhiForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
             else:
                 tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
-                model = LlavaLlamaForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
+                model = LlavaMistralForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
     else:
         # Load language model
         if model_base is not None:
@@ -138,14 +143,13 @@ def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, l
         vision_tower = model.get_vision_tower()
         if not vision_tower.is_loaded:
             vision_tower.load_model()
-        
-
-        vision_tower.to(device=model.device, dtype=torch.float16)
+        vision_tower.to(device='cuda', dtype=torch.float16)
         image_processor = vision_tower.image_processor
 
     if hasattr(model.config, "max_sequence_length"):
         context_len = model.config.max_sequence_length
     else:
         context_len = 2048
+    print("model loaded", model)
 
     return tokenizer, model, image_processor, context_len

llava/model/language_model/llava_llama.py

@@ -28,7 +28,7 @@ from ..llava_arch import LlavaMetaModel, LlavaMetaForCausalLM
 
 
 class LlavaConfig(LlamaConfig):
-    model_type = "llava"
+    model_type = "bakllava"
 
 
 class LlavaLlamaModel(LlavaMetaModel, LlamaModel):
@@ -136,5 +136,5 @@ class LlavaLlamaForCausalLM(LlamaForCausalLM, LlavaMetaForCausalLM):
         )
         return model_inputs
 
-AutoConfig.register("llava", LlavaConfig)
+AutoConfig.register("bakllava", LlavaConfig)
 AutoModelForCausalLM.register(LlavaConfig, LlavaLlamaForCausalLM)

llava/model/multimodal_encoder/builder.py

@@ -5,7 +5,7 @@ from .clip_encoder import CLIPVisionTower
 def build_vision_tower(vision_tower_cfg, **kwargs):
     vision_tower = getattr(vision_tower_cfg, 'mm_vision_tower', getattr(vision_tower_cfg, 'vision_tower', None))
     is_absolute_path_exists = os.path.exists(vision_tower)
-    if is_absolute_path_exists or vision_tower.startswith("openai") or vision_tower.startswith("laion"):
+    if is_absolute_path_exists or vision_tower.startswith("openai") or vision_tower.startswith("laion") or vision_tower.startswith("apple"):
         return CLIPVisionTower(vision_tower, args=vision_tower_cfg, **kwargs)
 
     raise ValueError(f'Unknown vision tower: {vision_tower}')

llava/model/multimodal_encoder/clip_encoder.py

@@ -2,6 +2,30 @@ import torch
 import torch.nn as nn
 
 from transformers import CLIPVisionModel, CLIPImageProcessor, CLIPVisionConfig
+from huggingface_hub import hf_hub_download
+import json
+
+def get_open_clip_image_processor(model_name):
+    config_path = hf_hub_download(model_name, filename="open_clip_config.json")
+
+    with open(config_path, 'r') as f:
+        config = json.load(f)
+    image_size = config['model_cfg']['vision_cfg']['image_size']
+    image_mean = config['preprocess_cfg']['mean']
+    image_std = config['preprocess_cfg']['std']
+    size = {"shortest_edge": image_size}
+    crop_size = {
+      "height": image_size,
+      "width": image_size
+    }
+
+    return CLIPImageProcessor(
+            image_size=image_size, 
+            image_mean=image_mean, 
+            image_std=image_std,
+            crop_size=crop_size,
+            size=size
+            )
 
 
 class CLIPVisionTower(nn.Module):
@@ -13,15 +37,18 @@ class CLIPVisionTower(nn.Module):
         self.vision_tower_name = vision_tower
         self.select_layer = args.mm_vision_select_layer
         self.select_feature = getattr(args, 'mm_vision_select_feature', 'patch')
-
         if not delay_load:
             self.load_model()
         else:
             self.cfg_only = CLIPVisionConfig.from_pretrained(self.vision_tower_name)
 
     def load_model(self):
-        self.image_processor = CLIPImageProcessor.from_pretrained(self.vision_tower_name)
         self.vision_tower = CLIPVisionModel.from_pretrained(self.vision_tower_name)
+        if self.vision_tower_name.startswith("apple") or self.vision_tower_name.startswith("laion"):
+            self.image_processor = get_open_clip_image_processor(self.vision_tower_name)
+        else:
+            self.image_processor = CLIPImageProcessor.from_pretrained(self.vision_tower_name)
+
         self.vision_tower.requires_grad_(False)
 
         self.is_loaded = True

llava/train/train.py

@@ -20,7 +20,7 @@ from dataclasses import dataclass, field
 import json
 import logging
 import pathlib
-from typing import Dict, Optional, Sequence, List
+from typing import Dict, Optional, Sequence, List, Union
 
 import torch
 
@@ -35,7 +35,14 @@ from llava.model import *
 from llava.mm_utils import tokenizer_image_token
 
 from PIL import Image
-
+import webdataset as wds
+import io
+import deepspeed
+import time
+from deepspeed.accelerator import get_accelerator
+import zipfile
+from webdataset_utils import get_wds_data
+import math
 
 local_rank = None
 
@@ -62,17 +69,30 @@ class ModelArguments:
 
 @dataclass
 class DataArguments:
-    data_path: str = field(default=None,
+    dataset_type: str = "webdataset"
+    dataset_resampled: bool = False
+    lengths_path: Optional[str] = None
+    data_path: Union[List[str], str] = field(default=None,
                            metadata={"help": "Path to the training data."})
     lazy_preprocess: bool = False
     is_multimodal: bool = False
     image_folder: Optional[str] = field(default=None)
     image_aspect_ratio: str = 'square'
     image_grid_pinpoints: Optional[str] = field(default=None)
+    train_data_weights: Optional[List[str]] = None
+    # dataloader_num_workers: Optional[int] = None
+    # seed: int = 0
 
 
 @dataclass
 class TrainingArguments(transformers.TrainingArguments):
+    num_training_samples:int = field(default=None)
+    resume_from_checkpoint:bool = False
+    deepspeed_config: str = field(default=None)
+    lr: float = field(default=1e-3)
+    beta1: float = field(default=0.5)
+    beta2: float = field(default=0.999)
+    num_train_epochs: int = field(default=1)
     cache_dir: Optional[str] = field(default=None)
     optim: str = field(default="adamw_torch")
     remove_unused_columns: bool = field(default=False)
@@ -85,6 +105,11 @@ class TrainingArguments(transformers.TrainingArguments):
             "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
         },
     )
+    dispatch_batches: bool = field(default=None)
+    pin_memory: bool = field(default=False)
+    resume: Optional[str] = field(default=None)
+    
+    # train_num_samples: int = field()
     double_quant: bool = field(
         default=True,
         metadata={"help": "Compress the quantization statistics through double quantization."}
@@ -104,6 +129,8 @@ class TrainingArguments(transformers.TrainingArguments):
     lora_weight_path: str = ""
     lora_bias: str = "none"
     group_by_modality_length: bool = field(default=False)
+    token: str = None
+    train_mode: str = "visual_instruction" # ["visual_instruction", "language_pretraining"]
 
 
 def maybe_zero_3(param, ignore_status=False, name=None):
@@ -617,7 +644,6 @@ def preprocess(
             tokenized_lens = _tokenize_fn([header] + [s["value"] for s in source], tokenizer)["input_ids_lens"]
         speakers = [sentence["from"] for sentence in source]
         _mask_targets(target, tokenized_lens, speakers)
-
     return dict(input_ids=input_ids, labels=targets)
 
 
@@ -634,6 +660,8 @@ class LazySupervisedDataset(Dataset):
         self.tokenizer = tokenizer
         self.list_data_dict = list_data_dict
         self.data_args = data_args
+        self.zip_file = zipfile.ZipFile(self.data_args.image_folder, 'r')
+
 
     def __len__(self):
         return len(self.list_data_dict)
@@ -664,7 +692,12 @@ class LazySupervisedDataset(Dataset):
             image_file = self.list_data_dict[i]['image']
             image_folder = self.data_args.image_folder
             processor = self.data_args.image_processor
-            image = Image.open(os.path.join(image_folder, image_file)).convert('RGB')
+            while True:
+                try:
+                    image = Image.open(io.BytesIO(self.zip_file.read(image_file)))
+                    break
+                except:
+                    pass
             if self.data_args.image_aspect_ratio == 'pad':
                 def expand2square(pil_img, background_color):
                     width, height = pil_img.size
@@ -721,8 +754,11 @@ class DataCollatorForSupervisedDataset(object):
         labels = torch.nn.utils.rnn.pad_sequence(labels,
                                                  batch_first=True,
                                                  padding_value=IGNORE_INDEX)
+    
+        
         input_ids = input_ids[:, :self.tokenizer.model_max_length]
         labels = labels[:, :self.tokenizer.model_max_length]
+
         batch = dict(
             input_ids=input_ids,
             labels=labels,
@@ -738,6 +774,90 @@ class DataCollatorForSupervisedDataset(object):
 
         return batch
 
+class WdsProcessor:
+    def __init__(self, tokenizer, data_args):
+        self.data_args = data_args
+        self.tokenizer = tokenizer
+        # processor = self.data_args.image_processor
+    def expand2square(self, pil_img, background_color):
+        width, height = pil_img.size
+        if width == height:
+            return pil_img
+        elif width > height:
+            result = Image.new(pil_img.mode, (width, width), background_color)
+            result.paste(pil_img, (0, (width - height) // 2))
+            return result
+        else:
+            result = Image.new(pil_img.mode, (height, height), background_color)
+            result.paste(pil_img, ((height - width) // 2, 0))
+            return result
+
+    def preprocess_wds(self, data):
+        
+        image, sources = data
+        has_image = 'image' in sources
+        sources = [sources]
+        image_processor = self.data_args.image_processor
+        if has_image:
+            if self.data_args.image_aspect_ratio == 'pad':
+                    
+                image = self.expand2square(image, tuple(int(x*255) for x in image_processor.image_mean))
+                image = image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
+            else:
+                image = image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
+            sources = preprocess_multimodal(
+                copy.deepcopy([e["conversations"] for e in sources]),
+                self.data_args)
+        else:
+            sources = copy.deepcopy([e["conversations"] for e in sources])
+        
+        data_dict = preprocess(
+                sources,
+                self.tokenizer,
+                has_image=has_image)
+
+        data_dict = dict(input_ids=data_dict["input_ids"][0],
+                             labels=data_dict["labels"][0])
+
+
+        if has_image:
+            data_dict['image'] = image
+        elif self.data_args.is_multimodal:
+            # image does not exist in the data, but the model is multimodal
+            crop_size = self.data_args.image_processor.crop_size
+            data_dict['image'] = torch.zeros(3, crop_size['height'], crop_size['width'])
+        
+        return data_dict
+
+def get_wds_dataset(tokenizer, data_args, training_args):
+    visual_instruction = training_args.train_mode == "visual_instruction"
+
+    round_fn = math.ceil
+    
+    if data_args.lengths_path:
+        with open(data_args.lengths_path, 'r') as f:
+            lengths = json.load(f)['length_list']
+            num_samples = len(lengths)
+    elif training_args.num_training_samples:
+        num_samples = training_args.num_training_samples
+    global_batch_size = training_args.per_device_train_batch_size  * training_args.world_size
+    num_batches = round_fn(num_samples / global_batch_size)
+    num_workers = max(1, training_args.dataloader_num_workers)
+    num_worker_batches = round_fn(num_batches / num_workers)  # per dataloader worker
+    num_batches = num_worker_batches * num_workers
+    training_args.max_steps = num_batches
+    data_args.train_num_samples = training_args.num_training_samples
+    
+    data_args.tokenizer = tokenizer
+    data_args.dataloader_num_workers = training_args.dataloader_num_workers
+    data_args.batch_size = training_args.per_device_train_batch_size
+    data_args.world_size = training_args.world_size
+    wds_processor = WdsProcessor(tokenizer, data_args)
+    train_data = get_wds_data(data_args, is_train=True, wds_processor=wds_processor.preprocess_wds) 
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer) if visual_instruction else None 
+    return dict(train_dataset=train_data,
+                eval_dataset=None,
+                data_collator=data_collator)
 
 def make_supervised_data_module(tokenizer: transformers.PreTrainedTokenizer,
                                 data_args) -> Dict:
@@ -788,16 +908,30 @@ def train():
                 cache_dir=training_args.cache_dir,
                 **bnb_model_from_pretrained_args
             )
+        elif 'mistral' in model_args.model_name_or_path.lower():
+            model = LlavaMistralForCausalLM.from_pretrained(
+                model_args.model_name_or_path,
+                cache_dir=training_args.cache_dir,
+                **bnb_model_from_pretrained_args
+            )
+        elif 'phi' in model_args.model_name_or_path:
+            model = LlavaPhiForCausalLM.from_pretrained(
+                model_args.model_name_or_path,
+                cache_dir=training_args.cache_dir,
+                **bnb_model_from_pretrained_args
+            )
         else:
             model = LlavaLlamaForCausalLM.from_pretrained(
                 model_args.model_name_or_path,
                 cache_dir=training_args.cache_dir,
+                token=training_args.token,
                 **bnb_model_from_pretrained_args
             )
     else:
         model = transformers.LlamaForCausalLM.from_pretrained(
             model_args.model_name_or_path,
             cache_dir=training_args.cache_dir,
+            token=training_args.token,
             **bnb_model_from_pretrained_args
         )
     model.config.use_cache = False
@@ -915,8 +1049,24 @@ def train():
                     if training_args.bf16 and module.weight.dtype == torch.float32:
                         module = module.to(torch.bfloat16)
 
-    data_module = make_supervised_data_module(tokenizer=tokenizer,
-                                              data_args=data_args)
+    if data_args.dataset_type == "webdataset":
+    
+        training_args.group_by_length = False     
+        data_module = get_wds_dataset(
+            tokenizer=tokenizer, 
+            data_args=data_args, 
+            training_args=training_args
+                )
+
+        
+    elif data_args.dataset_type == "files":
+        data_module = make_supervised_data_module(
+            tokenizer=tokenizer,
+            data_args=data_args
+            )
+    else:
+        ValueError(f"Unknown dataset type {data_args.dataset_type}! Dataset type should be euther `webdataset` or `files`")
+    
     trainer = LLaVATrainer(model=model,
                     tokenizer=tokenizer,
                     args=training_args,

llava/train/train_mem.py

@@ -3,9 +3,9 @@
 # Make it more memory efficient by monkey patching the LLaMA model with FlashAttn.
 
 # Need to call this before importing transformers.
-from llava.train.llama_flash_attn_monkey_patch import replace_llama_attn_with_flash_attn
+# from llava.train.llama_flash_attn_monkey_patch import replace_llama_attn_with_flash_attn
 
-replace_llama_attn_with_flash_attn()
+# replace_llama_attn_with_flash_attn()
 
 from llava.train.train import train
 

requirements.txt

@@ -16,7 +16,7 @@ shortuuid
 httpx==0.24.0
 deepspeed==0.9.5
 peft==0.4.0
-transformers==4.31.0
+transformers==4.36.0
 accelerate==0.21.0
 bitsandbytes==0.41.0
 scikit-learn==1.2.2