Update app.py

app.py

@@ -82,32 +82,25 @@ def model_generate_ans(img=None,img_audio=None,val_q=None):
             val_q_embeds    = merged_model.model.embed_tokens(val_q_tokenised).unsqueeze(0)
             val_combined_embeds.append(val_q_embeds)
 
-        # val_combined_embeds = []
-        # if img is not None:
-        #     #val_combined_embeds = torch.cat([val_combined_embeds, val_image_embeds, img_token_embeds], dim=1)
-        #     val_combined_embeds.append(val_image_embeds)
-        #     val_combined_embeds.append(img_token_embeds)
-        # if img_audio is not None:
-        #     #val_combined_embeds = torch.cat([val_combined_embeds, audio_embeds], dim=1)
-        #     val_combined_embeds.append(audio_embeds)
-        # if len(val_q) != 0:
-        #     #val_combined_embeds = torch.cat([val_combined_embeds, val_q_embeds], dim=1)
-        #     val_combined_embeds.append(val_q_embeds)
 
         val_combined_embeds = torch.cat(val_combined_embeds,dim=1)
-        
-        #val_combined_embeds = torch.cat([val_image_embeds, img_token_embeds, val_q_embeds], dim=1) # 4, 69, 2560
-        predicted_caption = torch.full((1,max_generate_length),50256).to(device)
+        predicted_caption = merged_model.generate(inputs_embeds=val_combined_embeds,
+                                                  max_new_tokens=max_generate_length,
+                                                  return_dict_in_generate = True)
+    
+        # predicted_caption = torch.full((1,max_generate_length),50256).to(device)
     
-        for g in range(max_generate_length):
-            phi_output_logits = merged_model(inputs_embeds=val_combined_embeds)['logits'] # 4, 69, 51200
-            predicted_word_token_logits = phi_output_logits[:, -1, :].unsqueeze(1) # 4,1,51200
-            predicted_word_token = torch.argmax(predicted_word_token_logits, dim = -1) # 4,1
-            predicted_caption[:,g] = predicted_word_token.view(1,-1)
-            next_token_embeds = phi_model.model.embed_tokens(predicted_word_token) # 4,1,2560
-            val_combined_embeds   = torch.cat([val_combined_embeds, next_token_embeds], dim=1)
+        # for g in range(max_generate_length):
+        #     phi_output_logits = merged_model(inputs_embeds=val_combined_embeds)['logits'] # 4, 69, 51200
+        #     predicted_word_token_logits = phi_output_logits[:, -1, :].unsqueeze(1) # 4,1,51200
+        #     predicted_word_token = torch.argmax(predicted_word_token_logits, dim = -1) # 4,1
+        #     predicted_caption[:,g] = predicted_word_token.view(1,-1)
+        #     next_token_embeds = phi_model.model.embed_tokens(predicted_word_token) # 4,1,2560
+        #     val_combined_embeds   = torch.cat([val_combined_embeds, next_token_embeds], dim=1)
             
-        predicted_captions_decoded = tokenizer.batch_decode(predicted_caption,ignore_index = 50256)[0]
+        #predicted_captions_decoded = tokenizer.batch_decode(predicted_caption,ignore_index = 50256)[0]
+        predicted_captions_decoded = tokenizer.batch_decode(predicted_caption.sequences[:, 1:])[0] 
+        predicted_captions_decoded = predicted_captions_decoded.replace("<|endoftext|>", "")  
     
     return predicted_captions_decoded