Spaces:
Running
on
Zero
Running
on
Zero
import os | |
import urllib.parse | |
import requests | |
import gradio as gr | |
from transformers import pipeline | |
class WikipediaBot: | |
SEARCH_TEMPLATE = "https://en.wikipedia.org/w/api.php?action=opensearch&search=%s&limit=1&namespace=0&format=json" | |
CONTENT_TEMPLATE = "https://en.wikipedia.org/w/api.php?format=json&action=query&prop=extracts&exintro&explaintext&redirects=1&titles=%s" | |
def __init__(self): | |
self.summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6") | |
def search_wikipedia(self, query): | |
query = urllib.parse.quote_plus(query) | |
data = requests.get(self.SEARCH_TEMPLATE % query).json() | |
if data and data[1]: | |
page = urllib.parse.quote_plus(data[1][0]) | |
content = requests.get(self.CONTENT_TEMPLATE % page).json() | |
content = list(content["query"]["pages"].values())[0]["extract"] | |
summary = self.summarizer(content, max_length=150, min_length=50, do_sample=False)[0]['summary_text'] | |
source = data[3][0] | |
return summary, source | |
else: | |
return "No results found.", "" | |
def chatbot_response(self, message, history): | |
summary, source = self.search_wikipedia(message) | |
response = f"Here's a summary of the top Wikipedia result:\n\n{summary}" | |
if source: | |
response += f"\n\nSource: {source}" | |
history.append((message, response)) | |
return history | |
def create_chatbot(): | |
return WikipediaBot() | |
demo = gr.Blocks() | |
with demo: | |
gr.Markdown("# Wikipedia Chatbot") | |
gr.Markdown("This chatbot queries the Wikipedia API and summarizes the top result.") | |
chatbot = gr.Chatbot() | |
msg = gr.Textbox(label="Enter your query") | |
clear = gr.Button("Clear") | |
bot = create_chatbot() | |
msg.submit(bot.chatbot_response, [msg, chatbot], chatbot).then( | |
lambda: gr.update(value=""), None, [msg], queue=False | |
) | |
clear.click(lambda: None, None, chatbot, queue=False) | |
if __name__ == "__main__": | |
demo.launch() |