Spaces:
Running
Running
import gradio as gr | |
from transformers import pipeline, RobertaTokenizer, RobertaForQuestionAnswering | |
import torch | |
# Load the model and tokenizer | |
model_name = "AventIQ-AI/roberta-chatbot" | |
tokenizer = RobertaTokenizer.from_pretrained(model_name) | |
model = RobertaForQuestionAnswering.from_pretrained(model_name) | |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
model = model.to(device) | |
# Initialize the question-answering pipeline | |
qa_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1) | |
# Define the function for the Gradio interface | |
def roberta_chatbot(context, question): | |
if not context or not question: | |
return "Please provide both context and a question." | |
# Get the model's answer | |
result = qa_pipeline(question=question, context=context) | |
answer = result.get('answer', 'Sorry, I could not find an answer.') | |
return answer | |
# Create the Gradio interface | |
iface = gr.Interface( | |
fn=roberta_chatbot, | |
inputs=[ | |
gr.Textbox(label="π Context", placeholder="Enter the context here...", lines=5), | |
gr.Textbox(label="β Question", placeholder="Enter your question here...", lines=2) | |
], | |
outputs=gr.Textbox(label="π€ Answer"), | |
title="π§ RoBERTa-Powered Chatbot", | |
description="Provide a context and ask a question. The RoBERTa-based chatbot will find the answer based on the given context.", | |
examples=[ | |
["Flight AI101 departs from New York at 10:00 AM and arrives in San Francisco at 1:30 PM. The flight duration is 5 hours and 30 minutes.", "What is the duration of Flight AI101?"], | |
["The Great Wall of China was built over several centuries to protect China's northern borders.", "Why was the Great Wall of China built?"] | |
], | |
theme="compact", | |
allow_flagging="never" | |
) | |
if __name__ == "__main__": | |
iface.launch() | |