Create app.py

app.py

@@ -0,0 +1,303 @@
+# ==========================================
+# 1. INITIAL SETUP & LIBRARIES
+# ==========================================
+import os
+import json
+import uuid
+import base64
+import whisper
+import pymupdf4llm
+import gradio as gr
+from datetime import datetime
+from huggingface_hub import InferenceClient
+from langchain_text_splitters import MarkdownTextSplitter
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.vectorstores import Chroma
+
+# ==========================================
+# 2. CONNECT TO AI APIS (Replaces Local Models)
+# ==========================================
+print("⏳ Connecting to Hugging Face APIs...")
+
+# Get token from environment variable (Set this in HF Spaces Secrets)
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+if not HF_TOKEN:
+    print("⚠️ WARNING: HF_TOKEN not found! The AI will not be able to generate responses.")
+
+# --- A. Mistral-7B API (The Writer/Scientist) ---
+text_client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3", token=HF_TOKEN)
+
+# --- B. Qwen2-VL API (The Eye) ---
+# We use the 7B version since the cloud API handles the compute!
+vision_client = InferenceClient("Qwen/Qwen2-VL-7B-Instruct", token=HF_TOKEN)
+
+# --- C. Local Embeddings & Whisper (Runs fine on CPU) ---
+embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+whisper_model = whisper.load_model("base")
+
+print("✅ APIs and Local Models Loaded Successfully!")
+
+# ==========================================
+# 3. GLOBAL STATE & HELPERS
+# ==========================================
+main_paper_retriever = None
+brainstorm_retriever = None
+main_extracted_images = []
+chat_history_file = "research_lab_history.json"
+
+if not os.path.exists(chat_history_file):
+    with open(chat_history_file, "w") as f: json.dump([], f)
+
+def save_to_json(user_msg, combined_ans, mode):
+    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    entry = {"timestamp": timestamp, "mode": mode, "user": user_msg, "assistant": combined_ans}
+    try:
+        with open(chat_history_file, "r") as f: history = json.load(f)
+    except: history = []
+    history.append(entry)
+    with open(chat_history_file, "w") as f: json.dump(history, f, indent=4)
+
+def process_pdf_to_markdown(pdf_path, extract_images=True):
+    """Converts PDF to Markdown. Optionally extracts images."""
+    global main_extracted_images
+    output_image_dir = "extracted_images"
+
+    if extract_images:
+        if os.path.exists(output_image_dir):
+            for f in os.listdir(output_image_dir): os.remove(os.path.join(output_image_dir, f))
+        else:
+            os.makedirs(output_image_dir, exist_ok=True)
+
+    try:
+        if extract_images:
+            md_text = pymupdf4llm.to_markdown(pdf_path, write_images=True, image_path=output_image_dir, image_format="png")
+            main_extracted_images = [os.path.join(output_image_dir, f) for f in os.listdir(output_image_dir) if f.endswith(('.png', '.jpg'))]
+            main_extracted_images.sort()
+        else:
+            md_text = pymupdf4llm.to_markdown(pdf_path, write_images=False)
+        return md_text
+    except Exception as e:
+        return ""
+
+# --- UPLOAD HANDLER 1: MAIN PAPER ---
+def process_main_paper(file_obj):
+    global main_paper_retriever
+    main_paper_retriever = None
+    if file_obj is None: return "⚠️ No file uploaded."
+
+    try:
+        unique_id = f"main_{uuid.uuid4().hex[:8]}"
+        md_content = process_pdf_to_markdown(file_obj.name, extract_images=True)
+        splitter = MarkdownTextSplitter(chunk_size=1000, chunk_overlap=200)
+        chunks = splitter.create_documents([md_content])
+        vectordb = Chroma.from_documents(documents=chunks, embedding=embedding_model, collection_name=unique_id)
+        main_paper_retriever = vectordb.as_retriever(search_kwargs={"k": 3})
+        return f"✅ Main Paper Ready!\n📘 Text: Indexed\n👁️ Images: {len(main_extracted_images)} Extracted"
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+
+# --- UPLOAD HANDLER 2: REFERENCE SET ---
+def process_brainstorm_papers(file_list):
+    global brainstorm_retriever
+    brainstorm_retriever = None
+    if not file_list: return "⚠️ No files uploaded."
+    if len(file_list) > 3: return "⚠️ Limit exceeded: Max 3 PDFs."
+
+    try:
+        combined_md = ""
+        names = []
+        for file_obj in file_list:
+            names.append(os.path.basename(file_obj.name))
+            text = process_pdf_to_markdown(file_obj.name, extract_images=False)
+            combined_md += f"\n\n--- PAPER: {os.path.basename(file_obj.name)} ---\n{text}\n"
+
+        unique_id = f"brainstorm_{uuid.uuid4().hex[:8]}"
+        splitter = MarkdownTextSplitter(chunk_size=1500, chunk_overlap=300)
+        chunks = splitter.create_documents([combined_md])
+        vectordb = Chroma.from_documents(documents=chunks, embedding=embedding_model, collection_name=unique_id)
+        brainstorm_retriever = vectordb.as_retriever(search_kwargs={"k": 5})
+        return f"✅ Knowledge Base Ready!\n📚 Papers: {', '.join(names)}"
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+
+def transcribe_audio(audio_path):
+    if audio_path is None: return ""
+    return whisper_model.transcribe(audio_path)["text"].strip()
+
+# ==========================================
+# 4. INTELLIGENCE LAYERS (API WRAPPERS)
+# ==========================================
+
+# Helper function to call Mistral API
+def ask_mistral(prompt):
+    try:
+        response = text_client.text_generation(prompt, max_new_tokens=1000, temperature=0.3)
+        return response
+    except Exception as e:
+        return f"⚠️ API Error (Mistral): {str(e)}"
+
+# Helper function to call Qwen API
+def ask_qwen(prompt, image_paths):
+    try:
+        messages = [{"role": "user", "content": []}]
+        for img_path in image_paths:
+            with open(img_path, "rb") as image_file:
+                b64_img = base64.b64encode(image_file.read()).decode('utf-8')
+            messages[0]["content"].append({
+                "type": "image_url",
+                "image_url": {"url": f"data:image/png;base64,{b64_img}"}
+            })
+        messages[0]["content"].append({"type": "text", "text": prompt})
+        
+        response = vision_client.chat_completion(messages=messages, max_tokens=150)
+        return response.choices[0].message.content
+    except Exception as e:
+        return f"⚠️ API Error (Qwen - Server might be busy): {str(e)}"
+
+# MODE 1: CHAT WITH MAIN PAPER
+def get_main_paper_response(question):
+    global main_paper_retriever, main_extracted_images
+    vision_context = ""
+
+    # Vision Pass
+    if main_extracted_images:
+        images_to_process = main_extracted_images[:3]
+        vision_prompt = f"Relate these images to: {question}"
+        vision_context = ask_qwen(vision_prompt, images_to_process)
+
+    # Text Pass
+    if main_paper_retriever:
+        docs = main_paper_retriever.invoke(question)
+        text_context = "\n\n".join(d.page_content for d in docs)
+        prompt = f"""[INST] Use the context to answer. Integrate visual insights if available.
+Markdown Context: {text_context}
+Visual Insight: {vision_context}
+Question: {question} [/INST]"""
+        return ask_mistral(prompt)
+    return "⚠️ Please upload Main Paper."
+
+# MODE 2: BRAINSTORM NOVELTY
+def get_novelty_response(question):
+    global brainstorm_retriever
+    if not brainstorm_retriever: return "⚠️ Upload Reference Papers."
+
+    docs = brainstorm_retriever.invoke(question)
+    context = "\n\n".join(d.page_content for d in docs)
+    prompt = f"""[INST] You are a Senior Research Scientist.
+Analyze these papers to find gaps and novelty.
+Context: {context}
+Task: Identify limitations in these methodologies and suggest a NOVEL approach or gap.
+Query: {question} [/INST]"""
+    return ask_mistral(prompt)
+
+# MODE 3: BRAINSTORM SETUP
+def get_setup_response(question):
+    global brainstorm_retriever
+    if not brainstorm_retriever: return "⚠️ Upload Reference Papers."
+
+    docs = brainstorm_retriever.invoke(question)
+    context = "\n\n".join(d.page_content for d in docs)
+    prompt = f"""[INST] You are a Research Architect.
+Based on the methodologies in the context, design a robust EXPERIMENTAL SETUP.
+Context: {context}
+Task: Propose Datasets, Evaluation Metrics, Baselines, and Hardware requirements to validate the proposed novelty.
+Query: {question} [/INST]"""
+    return ask_mistral(prompt)
+
+# MODE 4: GENERATE PAPER DRAFT
+def get_draft_response(question):
+    global brainstorm_retriever
+    if not brainstorm_retriever: return "⚠️ Upload Reference Papers."
+
+    docs = brainstorm_retriever.invoke(question)
+    context = "\n\n".join(d.page_content for d in docs)
+    prompt = f"""[INST] You are an Academic Writer.
+Write a structured Research Paper Draft (Abstract, Introduction, Methodology, Experiments).
+Use the context from reference papers to write the 'Related Work' section effectively.
+Context: {context}
+Task: Generate a draft for a paper about: {question} [/INST]"""
+    return ask_mistral(prompt)
+
+
+# ==========================================
+# 5. GRADIO UI
+# ==========================================
+def reset_chat(): return []
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🔬 AI Research Scientist Lab (Production Version)")
+    gr.Markdown("Pipeline: Analyze -> Find Novelty -> Design Setup -> Write Draft")
+
+    with gr.Row():
+        # --- LEFT: CONTROLS ---
+        with gr.Column(scale=1):
+
+            mode_radio = gr.Radio(
+                choices=[
+                    "1. Chat with Paper",
+                    "2. Brainstorm Novelty",
+                    "3. Brainstorm Setup",
+                    "4. Generate Paper Draft"
+                ],
+                value="1. Chat with Paper",
+                label="🧪 Research Stage"
+            )
+
+            gr.Markdown("---")
+            gr.Markdown("### 📂 Input Data")
+
+            file_main = gr.File(label="Target Paper (Stage 1)", file_types=[".pdf"])
+            status_main = gr.Textbox(label="Status", value="Waiting...", interactive=False)
+
+            file_refs = gr.File(label="Reference Papers (Stages 2-4)", file_types=[".pdf"], file_count="multiple")
+            status_refs = gr.Textbox(label="Status", value="Waiting...", interactive=False)
+
+            clear_btn = gr.Button("🗑️ Clear Workspace")
+
+        # --- RIGHT: WORKSPACE ---
+        with gr.Column(scale=2):
+            chatbot = gr.Chatbot(label="Lab Assistant", height=700)
+            audio_input = gr.Audio(sources=["microphone"], type="filepath", label="🎤 Dictate Idea")
+
+            with gr.Row():
+                msg_input = gr.Textbox(placeholder="Enter your query or research topic...", scale=4)
+                send_btn = gr.Button("🚀 Execute", variant="primary", scale=1)
+
+    # --- HANDLERS ---
+    file_main.change(fn=process_main_paper, inputs=file_main, outputs=status_main)
+    file_refs.change(fn=process_brainstorm_papers, inputs=file_refs, outputs=status_refs)
+    audio_input.stop_recording(fn=transcribe_audio, inputs=audio_input, outputs=msg_input)
+    clear_btn.click(fn=reset_chat, outputs=chatbot)
+
+    # --- MAIN ROUTER ---
+    def respond(message, history, mode):
+        if not message.strip(): return "", history
+        if history is None: history = []
+
+        # Route based on selected Stage
+        if mode == "1. Chat with Paper":
+            response = get_main_paper_response(message)
+        elif mode == "2. Brainstorm Novelty":
+            response = get_novelty_response(message)
+        elif mode == "3. Brainstorm Setup":
+            response = get_setup_response(message)
+        elif mode == "4. Generate Paper Draft":
+            response = get_draft_response(message)
+        else:
+            response = "Error: Unknown Mode"
+
+        # Log & Update
+        final_ans = f"**[{mode}]**\n{response}"
+        save_to_json(message, final_ans, mode)
+        history.append({"role": "user", "content": message})
+        history.append({"role": "assistant", "content": final_ans})
+
+        return "", history
+
+    msg_input.submit(respond, [msg_input, chatbot, mode_radio], [msg_input, chatbot])
+    send_btn.click(respond, [msg_input, chatbot, mode_radio], [msg_input, chatbot])
+
+print("🚀 Launching Production Research Scientist Lab...")
+# In HF Spaces, we don't need share=True
+demo.launch()