victor HF staff commited on
Commit
f43cbc3
·
1 Parent(s): fc2d69e

wikipedia chatbot

Browse files
Files changed (1) hide show
  1. app.py +53 -4
app.py CHANGED
@@ -1,7 +1,56 @@
 
 
 
1
  import gradio as gr
 
2
 
3
- def greet(name):
4
- return "Hello " + name + "!!"
 
5
 
6
- demo = gr.Interface(fn=greet, inputs="text", outputs="text")
7
- demo.launch()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import urllib.parse
3
+ import requests
4
  import gradio as gr
5
+ from transformers import pipeline
6
 
7
+ class WikipediaBot:
8
+ SEARCH_TEMPLATE = "https://en.wikipedia.org/w/api.php?action=opensearch&search=%s&limit=1&namespace=0&format=json"
9
+ CONTENT_TEMPLATE = "https://en.wikipedia.org/w/api.php?format=json&action=query&prop=extracts&exintro&explaintext&redirects=1&titles=%s"
10
 
11
+ def __init__(self):
12
+ self.summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
13
+
14
+ def search_wikipedia(self, query):
15
+ query = urllib.parse.quote_plus(query)
16
+ data = requests.get(self.SEARCH_TEMPLATE % query).json()
17
+ if data and data[1]:
18
+ page = urllib.parse.quote_plus(data[1][0])
19
+ content = requests.get(self.CONTENT_TEMPLATE % page).json()
20
+ content = list(content["query"]["pages"].values())[0]["extract"]
21
+ summary = self.summarizer(content, max_length=150, min_length=50, do_sample=False)[0]['summary_text']
22
+ source = data[3][0]
23
+ return summary, source
24
+ else:
25
+ return "No results found.", ""
26
+
27
+ def chatbot_response(self, message, history):
28
+ summary, source = self.search_wikipedia(message)
29
+ response = f"Here's a summary of the top Wikipedia result:\n\n{summary}"
30
+ if source:
31
+ response += f"\n\nSource: {source}"
32
+ history.append((message, response))
33
+ return history
34
+
35
+ def create_chatbot():
36
+ return WikipediaBot()
37
+
38
+ demo = gr.Blocks()
39
+
40
+ with demo:
41
+ gr.Markdown("# Wikipedia Chatbot")
42
+ gr.Markdown("This chatbot queries the Wikipedia API and summarizes the top result.")
43
+
44
+ chatbot = gr.Chatbot()
45
+ msg = gr.Textbox(label="Enter your query")
46
+ clear = gr.Button("Clear")
47
+
48
+ bot = create_chatbot()
49
+
50
+ msg.submit(bot.chatbot_response, [msg, chatbot], chatbot).then(
51
+ lambda: gr.update(value=""), None, [msg], queue=False
52
+ )
53
+ clear.click(lambda: None, None, chatbot, queue=False)
54
+
55
+ if __name__ == "__main__":
56
+ demo.launch()