Spaces:
Sleeping
Sleeping
| # ========================================== | |
| # 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 (OpenAI-Compatible) | |
| # ========================================== | |
| print("⏳ Connecting to Hugging Face APIs...") | |
| 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.") | |
| # Initialize a single, generic client | |
| # We do NOT bind the model name here to prevent the "text-generation" tagging error | |
| hf_client = InferenceClient(api_key=HF_TOKEN) | |
| # Local Embeddings & Whisper | |
| 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): | |
| 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 "" | |
| 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)}" | |
| 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 (STRICT CHAT ROUTING) | |
| # ========================================== | |
| def ask_mistral(prompt): | |
| try: | |
| messages = [{"role": "user", "content": prompt}] | |
| # This explicitly hits the /v1/chat/completions route | |
| response = hf_client.chat.completions.create( | |
| model="mistralai/Mistral-7B-Instruct-v0.2", | |
| messages=messages, | |
| max_tokens=1000, | |
| temperature=0.3 | |
| ) | |
| return response.choices[0].message.content | |
| except Exception as e: | |
| return f"⚠️ API Error (Mistral): {str(e)}" | |
| 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}) | |
| # Enforce chat completions route for Vision model too | |
| response = hf_client.chat.completions.create( | |
| model="Qwen/Qwen2-VL-7B-Instruct", | |
| messages=messages, | |
| max_tokens=150 | |
| ) | |
| return response.choices[0].message.content | |
| except Exception as e: | |
| return f"⚠️ API Error (Qwen - Server busy): {str(e)}" | |
| # MODE 1: CHAT WITH MAIN PAPER | |
| def get_main_paper_response(question): | |
| global main_paper_retriever, main_extracted_images | |
| vision_context = "" | |
| 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) | |
| 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(): | |
| 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") | |
| 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) | |
| 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) | |
| def respond(message, history, mode): | |
| if not message.strip(): return "", history | |
| if history is None: history = [] | |
| 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" | |
| 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...") | |
| demo.launch() |