victor's picture
victor HF staff
wikipedia chatbot
f43cbc3
raw
history blame
2.03 kB
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()