Update app.py

app.py

@@ -31,33 +31,36 @@ git_model_base.to(device)
 # vilt_model.to(device)
 
 def generate_answer_git(processor, model, image, question):
-    # prepare image
-    pixel_values = processor(images=image, return_tensors="pt").pixel_values
+    # # prepare image
+    # pixel_values = processor(images=image, return_tensors="pt").pixel_values
 
-    # prepare question
-    input_ids = processor(text=question, add_special_tokens=False).input_ids
-    input_ids = [processor.tokenizer.cls_token_id] + input_ids
-    input_ids = torch.tensor(input_ids).unsqueeze(0)
+    # # prepare question
+    # input_ids = processor(text=question, add_special_tokens=False).input_ids
+    # input_ids = [processor.tokenizer.cls_token_id] + input_ids
+    # input_ids = torch.tensor(input_ids).unsqueeze(0)
     
-    generated_ids = model.generate(pixel_values=pixel_values, input_ids=input_ids, max_length=50, return_dict_in_generate=True, output_scores=True)
-    print('scores:')
-    print(generated_ids.scores)
-    # scoresList0 = torch.softmax(generated_ids.scores[0], dim=1).flatten().tolist()
-    # print(scoresList0)
-    # scoresList1 = torch.softmax(generated_ids.scores[1], dim=1).flatten().tolist()
-    # print(scoresList1)
-    idx = generated_ids.scores[0].argmax(-1).item()
-    idx1 = generated_ids.scores[1].argmax(-1).item()
-    print(idx, idx1)
-    print(model.config.id2label)
-    ans = model.config.id2label[idx]
-    ans1 = model.config.id2label[idx1]
-    print(ans, ans1)
-    print('sequences:')
-    print(generated_ids.sequences)
-    print(generated_ids)
-    generated_answer = processor.batch_decode(generated_ids.sequences, skip_special_tokens=True)
-    print(generated_answer)
+    # generated_ids = model.generate(pixel_values=pixel_values, input_ids=input_ids, max_length=50, return_dict_in_generate=True, output_scores=True)
+    # print('scores:')
+    # print(generated_ids.scores)
+    # idx = generated_ids.scores[0].argmax(-1).item()
+    # idx1 = generated_ids.scores[1].argmax(-1).item()
+    # print(idx, idx1)
+    # print(model.config.id2label)
+    # ans = model.config.id2label[idx]
+    # ans1 = model.config.id2label[idx1]
+    # print(ans, ans1)
+    # print('sequences:')
+    # print(generated_ids.sequences)
+    # print(generated_ids)
+    # generated_answer = processor.batch_decode(generated_ids.sequences, skip_special_tokens=True)
+    # print(generated_answer)
+
+    encoding = processor(images=image, text=question, return_tensors="pt")
+
+    with torch.no_grad():
+        outputs = model(**encoding)
+    predicted_class_idx = outputs.logits.argmax(-1).item()
+    return model.config.id2label[predicted_class_idx]
     
    
     return 'haha'