Upload app.py

app.py

@@ -0,0 +1,215 @@
+import time
+from pathlib import Path
+import sys
+import os
+import logging
+
+# ==========================================
+# 0. Fix sys.path BEFORE importing anything from src/
+#    (src/generators/__init__.py etc. do bare "from generators.xxx import ..."
+#    imports that require the src/ folder itself to already be on sys.path)
+# ==========================================
+project_root = Path(__file__).parent.resolve()
+sys.path.insert(0, str(project_root))
+sys.path.insert(0, str(project_root / "src"))
+
+import gradio as gr
+import torch
+from src.config import load_config
+from src.models.load_model import ModelLoader
+from src.generators.program_generator import GenerationPipeline
+from src.generators.text_to_mongo_generator import TextToMongoGenerator
+from src.vectorDB.index_manager import IndexManager
+from src.rag.rag_pipeline import RAGPipeline
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# Environment variables to suppress warnings
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] = "0.0"
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
+# ==========================================
+# 1. Load Configuration & Models
+# ==========================================
+logger.info("Loading configuration...")
+config = load_config("configs/config.yaml")
+
+logger.info("Initializing Model Loader...")
+model_loader = ModelLoader(config)
+tokenizer = model_loader.load_tokenizer()
+base_model = model_loader.load_base_model()
+
+# Try to load the Fine-Tuned (LoRA) model if it exists
+lora_path = Path(config.lora.output_dir)
+if lora_path.exists() and (lora_path / "adapter_config.json").exists():
+    logger.info(f"✅ Found LoRA adapter at {lora_path}. Loading Fine-Tuned model...")
+    finetuned_model = model_loader.autodetect_saved_model(str(lora_path))
+else:
+    logger.warning("⚠️ No LoRA adapter found. Fine-Tuned model will fallback to Base model.")
+    finetuned_model = base_model
+
+# ==========================================
+# 2. Setup Generation Pipelines
+# ==========================================
+logger.info("Setting up generation pipelines...")
+base_pipeline = GenerationPipeline(base_model, tokenizer, config.generation)
+finetuned_pipeline = GenerationPipeline(finetuned_model, tokenizer, config.generation)
+
+base_generator = TextToMongoGenerator(base_pipeline, config)
+finetuned_generator = TextToMongoGenerator(finetuned_pipeline, config)
+
+# ==========================================
+# 3. Setup RAG Pipeline (Optional)
+# ==========================================
+rag_pipeline = None
+try:
+    index_manager = IndexManager(
+        embedding_model=config.embedding.external_model,
+        language=config.retrieval.structural_language,
+        semantic_weight=1.0 - config.retrieval.structural_weight_alpha,
+        ast_weight=config.retrieval.structural_weight_alpha,
+        save_dir="data/indices"
+    )
+
+    if index_manager._indices_exist():
+        index_manager.load_indices()
+        rag_pipeline = RAGPipeline(
+            retriever=index_manager,
+            generator=finetuned_pipeline,
+            config=config.generation
+        )
+        logger.info("✅ RAG Pipeline initialized successfully.")
+    else:
+        logger.warning("⚠️ RAG indices not found in 'data/indices'. RAG feature will be disabled.")
+except Exception as e:
+    logger.error(f"⚠️ Could not initialize RAG pipeline: {e}")
+
+
+# ==========================================
+# 4. Gradio Generation Logic
+# ==========================================
+def generate_mongo_query(question: str, model_choice: str, use_rag: bool, top_k: int):
+    if not question.strip():
+        return "Please enter a valid question.", "Latency: N/A"
+
+    start_time = time.time()
+
+    # RAG Generation
+    if use_rag and rag_pipeline is not None:
+        generated_query, context_metadata = rag_pipeline.generate_with_rag(question, top_k=int(top_k))
+        latency = (time.time() - start_time) * 1000
+
+        # Format context for display
+        context_str = "Retrieved Context:\n" + "-" * 30 + "\n"
+        for i, ctx in enumerate(context_metadata[:int(top_k)]):
+            q = ctx.get('question', 'N/A')
+            mql = ctx.get('generated_mongodb_query', ctx.get('mongodb_query', 'N/A'))
+            context_str += f"[{i + 1}] Q: {q}\n    MQL: {mql}\n\n"
+
+        return generated_query, f"⏱️ Latency: {latency:.2f} ms\n\n{context_str}"
+
+    # Standard Generation (Base or Fine-Tuned)
+    else:
+        if model_choice == "Fine-Tuned (LoRA)":
+            result = finetuned_generator.generate_single(question)
+        else:
+            result = base_generator.generate_single(question)
+
+        latency = result['latency_ms']
+        info_text = f"⏱️ Latency: {latency:.2f} ms\nModel: {model_choice}"
+        return result['generated_query'], info_text
+
+
+# ==========================================
+# 5. Build Gradio UI
+# ==========================================
+with gr.Blocks(title="Text-to-MongoDB Generator", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # 🚀 Text-to-MongoDB Query Generator
+        Convert natural language questions into executable MongoDB queries using **Base**, **Fine-Tuned**, or **RAG-enhanced** models.
+        """
+    )
+
+    with gr.Row():
+        # Left Column: Inputs
+        with gr.Column(scale=3):
+            question_input = gr.Textbox(
+                label="Natural Language Question",
+                placeholder="e.g., How many singers do we have?",
+                lines=2,
+                info="Type your question about the database here."
+            )
+
+            with gr.Row():
+                model_choice = gr.Radio(
+                    choices=["Base Model", "Fine-Tuned (LoRA)"],
+                    value="Fine-Tuned (LoRA)",
+                    label="Model Selection"
+                )
+                use_rag = gr.Checkbox(
+                    label="🧠 Use RAG (Retrieval-Augmented Generation)",
+                    value=False,
+                    info="Requires pre-built indices"
+                )
+
+            with gr.Accordion("⚙️ RAG Settings", open=False):
+                top_k_slider = gr.Slider(
+                    minimum=1, maximum=10, value=3, step=1,
+                    label="Top-K Retrieved Contexts",
+                    info="Number of similar examples to retrieve from the vector database."
+                )
+
+            generate_btn = gr.Button("Generate MongoDB Query", variant="primary", size="lg")
+
+        # Right Column: Outputs
+        with gr.Column(scale=2):
+            query_output = gr.Code(
+                label="Generated MongoDB Query",
+                language="javascript",
+                lines=6,
+                interactive=False
+            )
+            info_output = gr.Textbox(
+                label="Generation Info & Context",
+                lines=8,
+                interactive=False,
+                max_lines=15
+            )
+
+    gr.Markdown("---")
+    gr.Markdown("### 💡 Example Questions")
+    gr.Examples(
+        examples=[
+            ["How many singers do we have?"],
+            ["What is the average age of all singers from France?"],
+            ["Show the name and release year of the song by the youngest singer."],
+            ["Find the number of pets whose weight is heavier than 10."],
+            ["What are the different countries with singers above age 20?"],
+            ["List the name and date of each battle."],
+        ],
+        inputs=[question_input],
+        cache_examples=False
+    )
+
+    # Link button to function
+    generate_btn.click(
+        fn=generate_mongo_query,
+        inputs=[question_input, model_choice, use_rag, top_k_slider],
+        outputs=[query_output, info_output]
+    )
+
+# ==========================================
+# 6. Launch the App
+# ==========================================
+if __name__ == "__main__":
+    # Launch locally on port 7860
+    demo.launch(
+        server_name="0.0.0.0",  # Makes it accessible on your local network
+        server_port=7860,
+        share=False,  # Set to True if you want a public temporary URL
+        show_error=True
+    )