Create app.py

app.py

@@ -0,0 +1,234 @@
+import spaces  # must be imported before torch/other CUDA-touching modules (ZeroGPU requirement)
+import time
+from pathlib import Path
+import sys
+import os
+import logging
+
+# ==========================================
+# 0. Fix sys.path BEFORE importing anything from src/
+# ==========================================
+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 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.generators.doc_generator import DocGenerator
+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)
+
+# Text-to-MongoDB Generators
+base_generator = TextToMongoGenerator(base_pipeline, config)
+finetuned_generator = TextToMongoGenerator(finetuned_pipeline, config)
+
+# Documentation Generators
+base_doc_generator = DocGenerator(base_pipeline)
+finetuned_doc_generator = DocGenerator(finetuned_pipeline)
+
+# ==========================================
+# 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
+# ==========================================
+@spaces.GPU(duration=60)
+def generate_output(task, input_text, model_choice, use_rag, top_k):
+    if not input_text.strip():
+        return "Please enter a valid input.", "Latency: N/A"
+    
+    start_time = time.time()
+    
+    if task == "Text-to-MongoDB":
+        if use_rag and rag_pipeline is not None:
+            generated_query, context_metadata = rag_pipeline.generate_with_rag(input_text, top_k=int(top_k))
+            latency = (time.time() - start_time) * 1000
+            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}"
+        else:
+            if model_choice == "Fine-Tuned (LoRA)":
+                result = finetuned_generator.generate_single(input_text)
+            else:
+                result = base_generator.generate_single(input_text)
+            latency = result['latency_ms']
+            info_text = f"⏱️ Latency: {latency:.2f} ms\nModel: {model_choice}\nTask: Text-to-MongoDB"
+            return result['generated_query'], info_text
+            
+    elif task == "Documentation Generation":
+        if model_choice == "Fine-Tuned (LoRA)":
+            generated_doc = finetuned_doc_generator.generate_docstring(input_text)
+        else:
+            generated_doc = base_doc_generator.generate_docstring(input_text)
+        latency = (time.time() - start_time) * 1000
+        info_text = f"⏱️ Latency: {latency:.2f} ms\nModel: {model_choice}\nTask: Documentation Generation"
+        return generated_doc, info_text
+
+def update_ui(task):
+    """Dynamically updates input placeholders and output syntax highlighting based on the selected task."""
+    if task == "Text-to-MongoDB":
+        return (
+            gr.update(placeholder="e.g., How many singers do we have?", label="Natural Language Question"),
+            gr.update(language="javascript")
+        )
+    else:
+        return (
+            gr.update(placeholder="e.g., def add(a, b):\n    return a + b", label="Code Snippet"),
+            gr.update(language="python")
+        )
+
+# ==========================================
+# 5. Build Gradio UI
+# ==========================================
+with gr.Blocks(title="CodeGen", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # 🚀 CodeGen
+        Generate **MongoDB queries** or **Code Documentation** using Base, Fine-Tuned, or RAG-enhanced models.
+        """
+    )
+    
+    with gr.Row():
+        # Left Column: Controls
+        with gr.Column(scale=1):
+            task_selector = gr.Radio(
+                choices=["Text-to-MongoDB", "Documentation Generation"],
+                value="Text-to-MongoDB",
+                label="Select Task"
+            )
+            model_choice = gr.Radio(
+                choices=["Base Model", "Fine-Tuned (LoRA)"],
+                value="Fine-Tuned (LoRA)",
+                label="Model Selection"
+            )
+            
+            with gr.Accordion("⚙️ Advanced Settings", open=False):
+                use_rag = gr.Checkbox(
+                    label="🧠 Use RAG (Text-to-MongoDB only)", 
+                    value=False,
+                    info="Requires pre-built indices"
+                )
+                top_k_slider = gr.Slider(
+                    minimum=1, maximum=10, value=3, step=1, 
+                    label="Top-K Retrieved Contexts"
+                )
+            
+            generate_btn = gr.Button("Generate", variant="primary", size="lg")
+            
+        # Right Column: Inputs & Outputs
+        with gr.Column(scale=2):
+            input_box = gr.Textbox(
+                label="Natural Language Question",
+                placeholder="e.g., How many singers do we have?",
+                lines=6,
+                info="Type your question or paste your code snippet here."
+            )
+            output_box = gr.Code(
+                label="Generated Output", 
+                language="javascript", 
+                lines=10,
+                interactive=False
+            )
+            info_output = gr.Textbox(
+                label="Generation Info", 
+                lines=4, 
+                interactive=False
+            )
+
+    gr.Markdown("---")
+    gr.Markdown("### 💡 Example Prompts")
+    gr.Examples(
+        examples=[
+            ["How many singers do we have?", "Text-to-MongoDB"],
+            ["What is the average age of all singers from France?", "Text-to-MongoDB"],
+            ["def calculate_area(radius):\n    return 3.14159 * radius ** 2", "Documentation Generation"],
+            ["import os\ndef list_files(directory):\n    return os.listdir(directory)", "Documentation Generation"],
+        ],
+        inputs=[input_box, task_selector],
+        cache_examples=False
+    )
+
+    # Link events
+    task_selector.change(fn=update_ui, inputs=task_selector, outputs=[input_box, output_box])
+    
+    generate_btn.click(
+        fn=generate_output,
+        inputs=[task_selector, input_box, model_choice, use_rag, top_k_slider],
+        outputs=[output_box, info_output]
+    )
+
+# ==========================================
+# 6. Launch the App
+# ==========================================
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        show_error=True
+    )