import gradio as gr import requests import fitz import re import numpy as np import faiss from sentence_transformers import SentenceTransformer from groq import Groq from faster_whisper import WhisperModel # ========================= # INITIALIZE MODELS # ========================= embedding_model = SentenceTransformer("all-MiniLM-L6-v2") whisper_model = WhisperModel("base", compute_type="int8") # 🔑 PUT YOUR GROQ API KEY HERE client = Groq(api_key="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") # Use stable model MODEL_NAME = "llama-3.3-70b-versatile" # Global storage sections = {} section_texts = [] index = None # ========================= # PDF FUNCTIONS # ========================= def download_arxiv_pdf(arxiv_id): try: url = f"https://arxiv.org/pdf/{arxiv_id}.pdf" response = requests.get(url) response.raise_for_status() file_path = f"{arxiv_id}.pdf" with open(file_path, "wb") as f: f.write(response.content) return file_path except: return None def extract_text_from_pdf(pdf_path): doc = fitz.open(pdf_path) text = "" for page in doc: text += page.get_text() return text # ========================= # ROBUST SECTION EXTRACTION # ========================= def extract_sections(text): patterns = [ r"\n([IVX]+\.\s+[A-Z][A-Z\s]+)", # Roman r"\n(\d+\.\d+\.\d+\s+[^\n]+)", # 1.1.1 r"\n(\d+\.\d+\s+[^\n]+)", # 1.1 r"\n(\d+\.\s+[^\n]+)", # 1. r"\n(\d+\s+[^\n]+)", # 1 r"\n([A-Z][A-Z\s]{4,})\n" # ALL CAPS ] matches = [] for p in patterns: matches.extend(list(re.finditer(p, text))) matches = sorted(matches, key=lambda x: x.start()) extracted = {} for i, match in enumerate(matches): title = match.group(1).strip() start = match.end() end = matches[i+1].start() if i+1 < len(matches) else len(text) content = text[start:end].strip() if len(content) > 4000: content = content[:4000] extracted[title] = content # Add abstract manually abstract_match = re.search(r"Abstract(.*?)\n", text, re.DOTALL) if abstract_match: extracted["Abstract"] = abstract_match.group(1).strip() return extracted # ========================= # VECTOR STORE # ========================= def build_vector_store(sections_dict): global index, section_texts section_texts = list(sections_dict.values()) if len(section_texts) == 0: index = None return embeddings = embedding_model.encode(section_texts) embeddings = np.array(embeddings).astype("float32") dim = embeddings.shape[1] index = faiss.IndexFlatL2(dim) index.add(embeddings) # ========================= # LOAD PAPER # ========================= def load_paper(arxiv_id): global sections pdf_path = download_arxiv_pdf(arxiv_id) if pdf_path is None: return gr.update(choices=[]), "❌ Invalid arXiv ID" text = extract_text_from_pdf(pdf_path) sections = extract_sections(text) build_vector_store(sections) return gr.update(choices=list(sections.keys())), "✅ Paper Loaded Successfully" # ========================= # SUMMARY FUNCTION # ========================= def summarize_section(section_title): try: if not sections: return "❌ Load paper first" if section_title not in sections: return "❌ Section not found" content = sections[section_title] if not content: return "❌ Empty section" content = content[:4000] prompt = f""" Summarize this research section: - Main idea - Key concepts - Simple explanation - Importance Section: {section_title} Content: {content} """ response = client.chat.completions.create( model=MODEL_NAME, messages=[{"role": "user", "content": prompt}], temperature=0.3 ) return response.choices[0].message.content except Exception as e: return f"❌ Error:\n{str(e)}" # ========================= # RAG CHAT # ========================= def rag_chat(message, history): try: global index if index is None: history.append({"role": "assistant", "content": "❌ Load paper first"}) return history, "" query_embedding = embedding_model.encode([message]) query_embedding = np.array(query_embedding).astype("float32") D, I = index.search(query_embedding, k=3) retrieved = "\n\n".join([section_texts[i] for i in I[0]]) prompt = f""" Answer using ONLY this context. Context: {retrieved} Question: {message} """ response = client.chat.completions.create( model=MODEL_NAME, messages=[{"role": "user", "content": prompt}], temperature=0.2 ) answer = response.choices[0].message.content # ✅ FIXED FORMAT history.append({"role": "user", "content": message}) history.append({"role": "assistant", "content": answer}) return history, "" except Exception as e: history.append({"role": "assistant", "content": f"❌ Error:\n{str(e)}"}) return history, "" # ========================= # VOICE CHAT # ========================= def voice_chat(audio, history): try: if audio is None: return history, "" segments, _ = whisper_model.transcribe(audio) text = " ".join([seg.text for seg in segments]) return rag_chat(text, history) except Exception as e: history.append({"role": "assistant", "content": f"❌ Error:\n{str(e)}"}) return history, "" # ========================= # UI # ========================= with gr.Blocks() as demo: gr.Markdown("# 📚 ArXiv Research Assistant", elem_id="title") with gr.Row(): arxiv_input = gr.Textbox(label="Enter arXiv ID", scale=4) load_btn = gr.Button("Load Paper", variant="primary", scale=1) status = gr.Markdown() with gr.Row(): section_dropdown = gr.Dropdown(label="Sections", scale=3) summarize_btn = gr.Button("Generate Summary", variant="secondary", scale=1) summary_output = gr.Markdown() gr.Markdown("## 💬 Chat with Paper") chatbot = gr.Chatbot(height=400) with gr.Row(): msg = gr.Textbox(label="Ask a question", scale=4) send_btn = gr.Button("Send", variant="primary", scale=1) gr.Markdown("## 🎙 Voice Query") with gr.Row(): audio = gr.Audio(type="filepath", scale=4) voice_btn = gr.Button("Ask via Voice", scale=1) # Actions load_btn.click(load_paper, inputs=arxiv_input, outputs=[section_dropdown, status]) summarize_btn.click(summarize_section, inputs=section_dropdown, outputs=summary_output) send_btn.click(rag_chat, inputs=[msg, chatbot], outputs=[chatbot, msg]) voice_btn.click(voice_chat, inputs=[audio, chatbot], outputs=[chatbot, msg]) demo.launch( theme=gr.themes.Soft(), css=""" #title {text-align:center} """ )