Update app.py

app.py

@@ -73,19 +73,22 @@ def generate_answers(img=None, aud = None, q = None, max_tokens = 30):
     inputs_embeddings.append(end_iq_embeds)
     # Combine embeddings
     combined_embeds  = torch.cat(inputs_embeddings, dim=1)
-    print("---------",combined_embeds.shape)
-
-    for pos in range(max_tokens - 1):
-        model_output_logits = phi2_model.forward(inputs_embeds = combined_embeds)['logits']
-        print(model_output_logits.shape)
-        predicted_word_token_logits = model_output_logits[:, -1, :].unsqueeze(1)
-        predicted_word_token = torch.argmax(predicted_word_token_logits, dim = -1)
-        predicted_caption[:, pos] = predicted_word_token.view(1,-1).to('cpu')
-        print(predicted_caption)
-        next_token_embeds = phi2_model.model.embed_tokens(predicted_word_token)
-        combined_embeds   = torch.cat([combined_embeds, next_token_embeds], dim=1)
-        print("combined_embeds", combined_embeds.shape)
-    predicted_captions_decoded = tokenizer.batch_decode(predicted_caption,ignore_index = 50256)[0]
+    predicted_caption = phi2_model.generate(inputs_embeds=combined_embeds,
+                                                  max_new_tokens=max_tokens,
+                                                  return_dict_in_generate = True)
+
+    # for pos in range(max_tokens - 1):
+    #     model_output_logits = phi2_model.forward(inputs_embeds = combined_embeds)['logits']
+    #     print(model_output_logits.shape)
+    #     predicted_word_token_logits = model_output_logits[:, -1, :].unsqueeze(1)
+    #     predicted_word_token = torch.argmax(predicted_word_token_logits, dim = -1)
+    #     predicted_caption[:, pos] = predicted_word_token.view(1,-1).to('cpu')
+    #     print(predicted_caption)
+    #     next_token_embeds = phi2_model.model.embed_tokens(predicted_word_token)
+    #     combined_embeds   = torch.cat([combined_embeds, next_token_embeds], dim=1)
+    #     print("combined_embeds", combined_embeds.shape)
+    # 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