Update app.py

app.py

@@ -13,6 +13,7 @@ tokenizer  = AutoTokenizer.from_pretrained(phi_model_name, trust_remote_code=Tru
 processor  = AutoProcessor.from_pretrained(clip_model_name)
 tokenizer.pad_token = tokenizer.eos_token
 IMAGE_TOKEN_ID = 23893 # token for word comment
+QA_TOKEN_ID = 50295 # token for qa
 device = "cuda" if torch.cuda.is_available() else "cpu"
 clip_embed = 768
 phi_embed  = 2560
@@ -83,6 +84,12 @@ def model_generate_ans(img=None,img_audio=None,val_q=None):
             val_combined_embeds.append(val_q_embeds)
 
 
+        if img_audio is not None or len(val_q) != 0: # add QA Token
+            
+            QA_token_tensor = torch.tensor(QA_TOKEN_ID).to(device)
+            QA_token_embeds = merged_model.model.embed_tokens(QA_token_tensor).unsqueeze(0).unsqueeze(0)
+            val_combined_embeds.append(QA_token_embeds)
+            
         val_combined_embeds = torch.cat(val_combined_embeds,dim=1)
         predicted_caption = merged_model.generate(inputs_embeds=val_combined_embeds,
                                                   max_new_tokens=max_generate_length,