v1

app.py

@@ -1,5 +1,5 @@
 # A100 Zero GPU
-# import spaces
+import spaces
 
 # TroL Package
 import torch
@@ -33,10 +33,10 @@ question="What is the troll doing? Provide the detail in the image and imagine w
 model_1_8, tokenizer_1_8 = load_trol(link='TroL-1.8B')
 
 # loading model
-# model_3_8, tokenizer_3_8 = load_trol(link='TroL-3.8B')
+model_3_8, tokenizer_3_8 = load_trol(link='TroL-3.8B')
 
 # loading model
-# model_7, tokenizer_7 = load_trol(link='TroL-7B')
+model_7, tokenizer_7 = load_trol(link='TroL-7B')
 
 def threading_function(inputs, image_token_number, streamer, device, model, tokenizer, temperature, new_max_token, top_p):
 
@@ -55,7 +55,7 @@ def threading_function(inputs, image_token_number, streamer, device, model, toke
     generation_kwargs.update({'use_cache': True})
     return model.generate(**generation_kwargs)
 
-# @spaces.GPU
+@spaces.GPU
 def bot_streaming(message, history, link, temperature, new_max_token, top_p):
 
     # model selection
@@ -70,9 +70,9 @@ def bot_streaming(message, history, link, temperature, new_max_token, top_p):
         tokenizer = tokenizer_7
     
     # cpu -> gpu
-    # for param in model.parameters():
-    #     if not param.is_cuda:
-    #         param.data = param.to(accel.device)
+    for param in model.parameters():
+        if not param.is_cuda:
+            param.data = param.to(accel.device)
 
     # prompt type -> input prompt
     image_token_number = None
@@ -131,11 +131,11 @@ def bot_streaming(message, history, link, temperature, new_max_token, top_p):
     buffer = ""
     for character in response:
         buffer += character
-        time.sleep(0.015)
+        time.sleep(0.012)
         yield buffer
 
 demo = gr.ChatInterface(fn=bot_streaming,
-                        additional_inputs = [gr.Radio(["1.8B"], label="Size", info="Select one model size", value="1.8B"), gr.Slider(0, 1, 0.9, label="temperature"), gr.Slider(1, 1024, 128, label="new_max_token"), gr.Slider(0, 1, 0.95, label="top_p")],
+                        additional_inputs = [gr.Radio(["1.8B", "3.8B", "7B"], label="Size", info="Select one model size", value="7B"), gr.Slider(0, 1, 0.9, label="temperature"), gr.Slider(1, 1024, 128, label="new_max_token"), gr.Slider(0, 1, 0.95, label="top_p")],
                         additional_inputs_accordion="Generation Hyperparameters",
                         theme=gr.themes.Soft(),
                         title="TroL",

trol/arch_internlm2/modeling_internlm2.py

@@ -857,7 +857,7 @@ class InternLM2Model(InternLM2PreTrainedModel):
         self.vocab_size = config.vocab_size
         self.config = config
 
-        self.tok_embeddings = nn.Embedding(config.vocab_size,
+        self.tok_embeddings = nn.Embedding(config.vocab_size+1,
                                            config.hidden_size,
                                            self.padding_idx)
         self.layers = nn.ModuleList([

trol/arch_internlm2/modeling_trol.py

@@ -30,7 +30,7 @@ class TroLForCausalLM(InternLM2PreTrainedModel):
         # Model
         self.model = InternLM2Model(config)
         self.vocab_size = config.vocab_size
-        self.output = nn.Linear(config.hidden_size, config.vocab_size-1, bias=False)
+        self.output = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         self.max_length = config.max_length
 
         # Initialize weights and apply final processing

trol/load_trol.py

@@ -1,11 +1,16 @@
 import torch
 import warnings
 from config import *
-from peft import LoraConfig
 from transformers import BitsAndBytesConfig
 
 warnings.filterwarnings(action='ignore')
 
+def setting_trol_config(trol, tok_trol, image_special_token):
+    trol.config.image_token_index = tok_trol.convert_tokens_to_ids(image_special_token)
+    trol.config.ignore_index = -100
+    trol.config.pad_token_id = tok_trol.eos_token_id
+    trol.config.eos_token_id = tok_trol.eos_token_id
+
 def load_trol(link):
 
     """
@@ -16,21 +21,24 @@ def load_trol(link):
         from .arch_internlm2.tokenization_internlm2 import InternLM2Tokenizer as TroLTokenizer
         bits = 4
         path = TROL_1_8B
-        bit_quant_skip = ["vit", "vision_proj", "ffn", "output"]
+        image_special_token = "<image>"
+        bit_quant_skip = ["vit", "vision_proj", "ffn", "output", "trol_gating"]
 
     elif link == 'TroL-3.8B':
         from trol.arch_phi3.modeling_trol import TroLForCausalLM 
         from transformers import LlamaTokenizerFast as TroLTokenizer
         bits = 8
         path = TROL_3_8B
-        bit_quant_skip = ["vision_model", "vision_proj", "lm_head"]
+        image_special_token = "<IMG_CONTEXT>"
+        bit_quant_skip = ["vision_model", "vision_proj", "lm_head", "trol_gating"]
 
     elif link == 'TroL-7B':
         from .arch_internlm2.modeling_trol import TroLForCausalLM
         from .arch_internlm2.tokenization_internlm2 import InternLM2Tokenizer as TroLTokenizer
         bits = 4
         path = TROL_7B
-        bit_quant_skip = ["vit", "vision_proj", "ffn", "output"]
+        image_special_token = "<image>"
+        bit_quant_skip = ["vit", "vision_proj", "ffn", "output", "trol_gating"]
     else:
         raise Exception("Unsupported Link")
 
@@ -68,10 +76,16 @@ def load_trol(link):
     except:
         del huggingface_config["attn_implementation"]
         trol = TroLForCausalLM.from_pretrained(path, **huggingface_config)
+        trol.config.llm_config.use_cache = False
+
+    # setting config
+    setting_trol_config(trol, tok_trol, image_special_token)
+
 
-    # wrapping
+    # trol gating load
+    from huggingface_hub import hf_hub_download
     try:
-        trol = trol.cuda()
+        trol.model.trol_gating.load_state_dict(torch.load(hf_hub_download(repo_id=path, filename="trol_gating.pt")))
     except:
-        pass
+        trol.language_model.model.trol_gating.load_state_dict(torch.load(hf_hub_download(repo_id=path, filename="trol_gating.pt")))
     return trol, tok_trol