import gradio as gr
from transformers import pipeline, AutoModelForQuestionAnswering, AutoTokenizer

# Load the pre-trained model and tokenizer
model_name = "distilbert-base-uncased-distilled-squad"
model = AutoModelForQuestionAnswering.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Load custom knowledge document
with open("knowledge.txt", "r") as file:
    custom_knowledge = file.read()

# Initialize the question-answering pipeline
qa_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer)

def answer_question(question):
    result = qa_pipeline(question=question, context=custom_knowledge)
    return result['answer']

# Set up the Gradio interface
interface = gr.Interface(
    fn=answer_question,
    inputs="text",
    outputs="text",
    title="Custom Knowledge QA",
    description="Ask questions based on the custom knowledge document."
)

if __name__ == "__main__":
    interface.launch()