Update lavis/models/protein_models/protein_function_opt.py

lavis/models/protein_models/protein_function_opt.py

@@ -98,26 +98,15 @@ class Blip2ProteinMistral(Blip2ProteinBase):
 
         self.mistral_tokenizer = LlamaTokenizer.from_pretrained("teknium/OpenHermes-2.5-Mistral-7B")
         # self.mistral_tokenizer = LlamaTokenizer.from_pretrained("/cluster/home/wenkai/.cache/huggingface/hub/models--teknium--OpenHermes-2.5-Mistral-7B", use_fast=False)
-        # configuration = MistralConfig()
         self.mistral_tokenizer.pad_token = '<pad>'
-        self.mistral_model = MistralForCausalLM.from_pretrained("teknium/OpenHermes-2.5-Mistral-7B", torch_dtype=torch.float16)
-        # self.mistral_model = MistralForCausalLM.from_pretrained("/cluster/home/wenkai/.cache/huggingface/hub/models--teknium--OpenHermes-2.5-Mistral-7B", torch_dtype=torch.float16)
-        # self.mistral_model = MistralForCausalLM(configuration)
-        for name, param in self.mistral_model.named_parameters():
-            param.requires_grad = False
-        #self.mistral_model.lm_head = self.mistral_model.lm_head.float()
-        #for param in self.mistral_model.lm_head.parameters():
-        #    param.requires_grad = True
-
-        #self.eos_token_id = self.mistral_tokenizer(
-        #    "\n", add_special_tokens=False
-        #).input_ids[0]
+        # self.mistral_model = MistralForCausalLM.from_pretrained("teknium/OpenHermes-2.5-Mistral-7B", torch_dtype=torch.float16)
+        self.mistral_model = None
+
         self.eos_token_id = self.mistral_tokenizer(
             "\n", add_special_tokens=False
         ).input_ids[1]
-        print(f"LLM hidden size: {self.mistral_model.config.hidden_size}")
         self.opt_proj = nn.Linear(
-            self.Qformer.config.hidden_size, self.mistral_model.config.hidden_size
+            self.Qformer.config.hidden_size, 4096
         )
 
         self.max_txt_len = max_txt_len
@@ -191,7 +180,6 @@ class Blip2ProteinMistral(Blip2ProteinBase):
         )
         targets = torch.cat([empty_targets, targets], dim=1)
 
-        #inputs_embeds = self.mistral_model.model.decoder.embed_tokens(mistral_tokens.input_ids)
         inputs_embeds = self.mistral_model.model.embed_tokens(mistral_tokens.input_ids)
         inputs_embeds = torch.cat([inputs_mistral, inputs_embeds], dim=1)
         attention_mask = torch.cat([atts_mistral, mistral_tokens.attention_mask], dim=1)
@@ -209,6 +197,7 @@ class Blip2ProteinMistral(Blip2ProteinBase):
     @torch.no_grad()
     def generate(
             self,
+            mistral_model,
             samples,
             # use_nucleus_sampling=False,
             num_beams=15,
@@ -262,8 +251,8 @@ class Blip2ProteinMistral(Blip2ProteinBase):
                 truncation=True,
                 max_length=self.max_txt_len,
             ).to(self.device)
-            # inputs_embeds = self.mistral_model.model.decoder.embed_tokens(mistral_tokens.input_ids)
-            inputs_embeds = self.mistral_model.model.embed_tokens(mistral_tokens.input_ids)
+            
+            inputs_embeds = mistral_model.model.embed_tokens(mistral_tokens.input_ids)
             inputs_embeds = torch.cat([inputs_mistral, inputs_embeds], dim=1)
             attention_mask = torch.cat([atts_mistral, mistral_tokens.attention_mask], dim=1)
             # if name[0] == 'Pin':
@@ -275,7 +264,7 @@ class Blip2ProteinMistral(Blip2ProteinBase):
             #num_txt = 15
             #return_num_txt = 10
             with torch.no_grad():
-                outputs = self.mistral_model.generate(inputs_embeds=inputs_embeds, attention_mask=attention_mask, min_length=min_length,
+                outputs = mistral_model.generate(inputs_embeds=inputs_embeds, attention_mask=attention_mask, min_length=min_length,
                                                   max_new_tokens=max_length, temperature=temperature, return_dict_in_generate=True,
                                                   output_scores=True,
                                                   repetition_penalty=repetition_penalty, num_beams=num_beams,