CannaAssist / app.py
CannaTech's picture
Update app.py
fffbd16 verified
raw
history blame
963 Bytes
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()