#허깅페이스에서 돌아갈 수 있도록 바꾸어 보았음

import gradio as gr
import torch
from transformers import BertTokenizer, BertForSequenceClassification

# 모델 초기화 및 가중치 불러오기
model_name = 'microsoft/git-base-vqav2'  # 사용할 모델의 이름
model = BertForSequenceClassification.from_pretrained(model_name)
tokenizer = BertTokenizer.from_pretrained(model_name)

# 예측 함수 정의
def predict_answer(image, question):
    inputs = tokenizer(question, return_tensors='pt')
    input_ids = inputs['input_ids']
    attention_mask = inputs['attention_mask']

    # 이미지와 관련된 처리 수행
    # 이미지 처리 코드를 여기에 추가해야 합니다 (입력된 이미지에 대한 전처리 등)

    # 모델에 입력 데이터를 전달하여 예측 수행
    with torch.no_grad():
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)

    # 예측 결과에서 가장 높은 확률을 가진 레이블 ID 가져오기
    predicted_label_id = torch.argmax(outputs.logits).item()
    predicted_label = id_to_label_fn(predicted_label_id)

    return predicted_label

iface = gr.Interface(
    fn=predict_answer,
    inputs=["image", "text"],
    outputs="text",
    title="Visual Question Answering",
    description="Input an image and a question to get the model's answer.",
    example=[
        "https://your_image_url.jpg",
        "What is shown in the image?"
    ]
)

iface.launch()