edit code for nle inference

app.py

@@ -1,13 +1,63 @@
-import gradio as gr
+import os
+import csv
+import json
 import torch
+import argparse
+import pandas as pd
+import torch.nn as nn
+from tqdm import tqdm
+from collections import defaultdict
+from transformers.models.llama.tokenization_llama import LlamaTokenizer
+from torch.utils.data import DataLoader
+from mplug_owl_video.modeling_mplug_owl import MplugOwlForConditionalGeneration
+from mplug_owl_video.processing_mplug_owl import MplugOwlImageProcessor, MplugOwlProcessor
+from peft import LoraConfig, get_peft_model
+from data_utils.xgpt3_dataset import MultiModalDataset
+from utils import batchify
+
+import gradio as gr
+from entailment_inference import get_scores
+
 print(f"Is CUDA available: {torch.cuda.is_available()}")
 # True
 print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
 # Tesla T4
 
-# def greet(name):
-#     return "Hello " + name + "!!"
+tokenizer = LlamaTokenizer.from_pretrained(pretrained_ckpt)
+image_processor = MplugOwlImageProcessor.from_pretrained(pretrained_ckpt)
+processor = MplugOwlProcessor(image_processor, tokenizer)
+
+
+# Instantiate model
+model = MplugOwlForConditionalGeneration.from_pretrained(
+    pretrained_ckpt,
+    torch_dtype=torch.bfloat16,
+    device_map={'':0}
+)
+
+for name, param in model.named_parameters():
+    param.requires_grad = False
+peft_config = LoraConfig(
+    target_modules=r'.*language_model.*\.(q_proj|v_proj|k_proj|o_proj|gate_proj|down_proj|up_proj)', 
+    inference_mode=True, 
+    r=32, 
+    lora_alpha=16, 
+    lora_dropout=0.05
+)
+model = get_peft_model(model, peft_config)
+model.print_trainable_parameters()
+with open(trained_ckpt, 'rb') as f:
+    ckpt = torch.load(f, map_location = torch.device(f"cuda:0"))
+model.load_state_dict(ckpt)
+model = model.to(torch.bfloat16)
+print('Model Loaded')
 
-# iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-# iface.launch()
+PROMPT = """The following is a conversation between a curious human and AI assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.
+Human: <|video|>
+Human: Does this video entail the description: ""A basketball team walking off the field while the audience claps.""?
+AI: """
 
+valid_data = MultiModalDataset("examples/y5xuvHpDPZQ_000005_000015.mp4", PROMPT, tokenizer, processor, max_length = 256, loss_objective = 'sequential')
+dataloader = DataLoader(valid_data, pin_memory=True, collate_fn=batchify)
+score = get_scores(model, tokenizer, dataloader)
+print(score)

data_utils/xgpt3_dataset.py

@@ -36,37 +36,32 @@ def load_jsonl(filename):
 class MultiModalDataset(Dataset):
     """MultiModal dataset"""
 
-    def __init__(self, input_file, tokenizer, processor,
+    def __init__(self, videopath, text, tokenizer, processor,
                  max_length=2048,
                  media_tokens=['<image>', '<|video|>'], loss_objective = 'sequential'):
         
         args = get_args()
 
         self.loss_objective = loss_objective
-        if 'sequential' in self.loss_objective:
-            self.dataset = pd.read_csv(input_file) 
-            self.dataset = self.dataset.dropna()
-        else:
+        if 'sequential' not in self.loss_objective:
             raise NotImplementedError('dataset loader not implemented for other loss objectives')
         
-        self.dataset = pd.read_csv(input_file)
+        self.videopath = videopath
+        self.text = text
         self.tokenizer = tokenizer
         self.max_length = max_length
         self.processor = processor
         self.media_tokens = {k: -int(i+1) for i, k in enumerate(media_tokens)}
         self.media_lengths = {'<image>': 1+64,'<|video|>': 1+64}
         print("num_media_token: ", self.media_lengths)
-        print(len(self.dataset))
         self.bucket = {}
 
     def __len__(self):
-        return len(self.dataset)
+        return 1
 
     def __getitem__(self, index):
-        
-        data = self.dataset.iloc[index]
-        videopath = data['videopath']
-        caption   = data['caption']
+        videopath = self.videopath
+        caption   = self.text
         video_input = self.processor(videos=[videopath], num_frames=32, return_tensors='pt') # video_pixel_values
         text_input = self._extract_text_token_from_conversation(caption, self.max_length, index)
         item  = {'video': video_input, 'text': text_input, 'videopath': videopath, 'caption': caption}

entailment_inference.py

@@ -15,18 +15,7 @@ from peft import LoraConfig, get_peft_model
 from data_utils.xgpt3_dataset import MultiModalDataset
 from utils import batchify
 
-parser = argparse.ArgumentParser()
 
-parser.add_argument('--input_csv', type = str, required = True, help = 'input json file')
-parser.add_argument('--output_csv', type = str, help = 'output csv with scores')
-parser.add_argument('--pretrained_ckpt', type = str, required = True, help = 'pretrained ckpt')
-parser.add_argument('--trained_ckpt', type = str, help = 'trained ckpt')
-parser.add_argument('--lora_r', type = int, default = 32)
-parser.add_argument('--use_lora', action = 'store_true', help = 'lora model')
-parser.add_argument('--all-params', action = 'store_true', help = 'use all params of the model')
-parser.add_argument('--batch_size', type = int, default = 32)
-
-args = parser.parse_args()
 softmax = nn.Softmax(dim=2)
 
 def get_entail(logits, input_ids, tokenizer):
@@ -47,7 +36,6 @@ def get_entail(logits, input_ids, tokenizer):
     return entailment
 
 def get_scores(model, tokenizer, dataloader):
-
     with torch.no_grad():
         for index, inputs in tqdm(enumerate(dataloader)):
             for k, v in inputs.items():
@@ -60,63 +48,4 @@ def get_scores(model, tokenizer, dataloader):
                                 non_media_mask = inputs['non_media_mask'], prompt_mask = inputs['prompt_mask'])
             logits = outputs['logits']
             entail_scores = get_entail(logits, inputs['input_ids'], tokenizer)
-            for m in range(len(entail_scores)):
-                with open(args.output_csv, 'a') as f:
-                    writer = csv.writer(f)
-                    writer.writerow([inputs['videopaths'][m], inputs['captions'][m], entail_scores[m].item()])
-            print(f"Batch {index} Done")
-
-def main():
-    
-    pretrained_ckpt = args.pretrained_ckpt
-
-    # Processors
-    tokenizer = LlamaTokenizer.from_pretrained(pretrained_ckpt)
-    image_processor = MplugOwlImageProcessor.from_pretrained(pretrained_ckpt)
-    processor = MplugOwlProcessor(image_processor, tokenizer)
-
-    valid_data = MultiModalDataset(args.input_csv, tokenizer, processor, max_length = 256, loss_objective = 'sequential')
-    dataloader = DataLoader(valid_data, batch_size=args.batch_size, pin_memory=True, collate_fn=batchify)
-    
-    # Instantiate model
-    model = MplugOwlForConditionalGeneration.from_pretrained(
-        pretrained_ckpt,
-        torch_dtype=torch.bfloat16,
-        device_map={'':0}
-    )
-
-    if args.use_lora:
-        for name, param in model.named_parameters():
-            param.requires_grad = False
-        if args.all_params:
-            peft_config = LoraConfig(
-                target_modules=r'.*language_model.*\.(q_proj|v_proj|k_proj|o_proj|gate_proj|down_proj|up_proj)', 
-                inference_mode=True, 
-                r=args.lora_r, 
-                lora_alpha=16, 
-                lora_dropout=0.05
-            )
-        else:
-            peft_config = LoraConfig(
-                target_modules=r'.*language_model.*\.(q_proj|v_proj|k_proj|o_proj)', 
-                inference_mode=True, 
-                r=args.lora_r, 
-                lora_alpha=16, 
-                lora_dropout=0.05
-            )
-            
-        model = get_peft_model(model, peft_config)
-        model.print_trainable_parameters()
-
-        with open(args.trained_ckpt, 'rb') as f:
-            ckpt = torch.load(f, map_location = torch.device(f"cuda:0"))
-        model.load_state_dict(ckpt)
-        model = model.to(torch.bfloat16)
-        print('Model Loaded')
-
-    model.eval()
-
-    get_scores(model, tokenizer, dataloader)
-
-if __name__  == "__main__":
-    main()
+            return entail_scores[0].item()