from models.blip_model import blip_answer
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
from config import DEVICE

tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
model = AutoModelForSeq2SeqLM.from_pretrained(
    "google/flan-t5-base"
).to(DEVICE)

model.eval()

def reasoning_answer(image, question):

    caption, base_answer = blip_answer(image, question)

    prompt = f"""
Scene:
{caption}

Question:
{question}

Base Answer:
{base_answer}

Provide:
Final Answer:
Explanation:
"""

    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)

    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=100)

    text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return caption, base_answer, text