search / app.py
HedronCreeper's picture
Update app.py
ed52d6a verified
import gradio as gr
import requests
import re
from transformers import pipeline
# Loading the model
summarizer = pipeline(
"text-generation",
model="HuggingFaceTB/SmolLM2-135M-Instruct",
device=-1
)
def clean_query(query):
query = query.lower()
query = re.sub(r"what is|who is|tell me about|explain|define", "", query)
query = re.sub(r'[^\w\s]', '', query)
return query.strip()
def get_wikipedia_data(query):
try:
cleaned = clean_query(query)
if not cleaned:
return None, []
search_url = f"https://en.wikipedia.org/api/rest_v1/page/summary/{cleaned.replace(' ', '_')}"
headers = {'User-Agent': 'MiniWikiBot/1.0 (contact@example.com)'}
resp = requests.get(search_url, headers=headers, timeout=5)
if resp.status_code == 200:
data = resp.json()
return data.get("extract"), []
suggest_url = f"https://en.wikipedia.org/w/api.php?action=opensearch&search={cleaned}&limit=3&namespace=0&format=json"
s_resp = requests.get(suggest_url, headers=headers, timeout=5)
suggestions = s_resp.json()[1] if s_resp.status_code == 200 else []
return None, suggestions
except:
return None, []
def ai_summarize(text, progress=gr.Progress()):
if not text or len(text) < 10 or text.startswith("❌"):
return "⚠️ Please search for a valid topic first."
try:
progress(0.1, desc="Initializing AI...")
prompt = f"<|im_start|>user\nSummarize the important info from this text:\n{text}\n<|im_end|>\n<|im_start|>assistant\n"
progress(0.3, desc="Generating summary...")
result = summarizer(
prompt,
max_new_tokens=200,
do_sample=True,
temperature=0.7,
top_p=0.9,
return_full_text=False
)
progress(0.9, desc="Cleaning up...")
summary = result[0]['generated_text'].strip()
summary = re.split(r'<\|im_start\|>|<\|im_end\|>', summary)[0].strip()
if not summary:
return "The AI was unable to generate a summary. Try a different topic!"
return summary
except Exception as e:
return f"❌ AI Error: {str(e)}"
def process_search(query):
text, sugs = get_wikipedia_data(query)
if text:
# Return Wikipedia text, clear AI output, and show AI button
return text, "", gr.update(visible=True)
if sugs:
sug_text = "❌ Not found. Did you mean:\n" + "\n".join([f"β€’ {s}" for s in sugs])
return sug_text, "", gr.update(visible=False)
return "❌ No results found.", "", gr.update(visible=False)
with gr.Blocks(analytics_enabled=False) as demo:
gr.Markdown("# πŸ€– Mini Wikipedia AI")
with gr.Row():
input_query = gr.Textbox(label="Ask something", placeholder="e.g. Black Holes", scale=4)
search_btn = gr.Button("Search", variant="primary", scale=1)
wiki_output = gr.Textbox(label="Wikipedia Summary", interactive=False, lines=8)
with gr.Row():
ai_btn = gr.Button("✨ Summarize with AI", visible=False)
ai_output = gr.Textbox(label="AI Summary Result", interactive=False, placeholder="AI summary will appear here...")
# Logic for searching
search_btn.click(process_search, inputs=[input_query], outputs=[wiki_output, ai_output, ai_btn])
input_query.submit(process_search, inputs=[input_query], outputs=[wiki_output, ai_output, ai_btn])
# Logic for AI Summarization
ai_btn.click(
fn=lambda: "",
outputs=[ai_output]
).then(
fn=ai_summarize,
inputs=[wiki_output],
outputs=[ai_output]
)
if __name__ == "__main__":
demo.launch(
theme=gr.themes.Soft(),
show_error=True
)