import gradio as gr
from transformers import RagTokenizer, RagTokenForGeneration

# Load tokenizer and model
tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
model = RagTokenForGeneration.from_pretrained("facebook/rag-token-nq")

def rag_generate(query):
    # Tokenize the input question
    inputs = tokenizer(query, return_tensors="pt")

    # Generate output
    generated_ids = model.generate(**inputs)

    # Decode the generated response
    answer = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    
    return answer

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 RAG Token QA with facebook/rag-token-nq")
    with gr.Row():
        question = gr.Textbox(label="Ask your question")
    with gr.Row():
        answer = gr.Textbox(label="Answer")

    submit_btn = gr.Button("Generate Answer")
    submit_btn.click(fn=rag_generate, inputs=question, outputs=answer)

demo.launch()