Blame for app.py

import spaces
import time
from pathlib import Path
import sys
import os
import logging

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
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 & Generators
# ==========================================
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 (linked to base & fine-tuned)
base_mongo_generator = TextToMongoGenerator(base_pipeline, config)
finetuned_mongo_generator = TextToMongoGenerator(finetuned_pipeline, config)

# Documentation Generators (linked to base & fine-tuned)
base_doc_generator = DocGenerator(base_pipeline)
finetuned_doc_generator = DocGenerator(finetuned_pipeline)

logger.info("✅ All generators initialized (MongoDB + Documentation) × (Base + Fine-Tuned)")

# ==========================================
# 3. Setup RAG Pipeline & Index Manager
# ==========================================
rag_pipeline = None
index_manager = 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. RAG Helper for Documentation
# ==========================================
def generate_doc_with_rag(code_snippet, pipeline, top_k=3):
    """Generate documentation using RAG-retrieved context."""
    if index_manager is None:
        # Fallback: no RAG available, use plain generation
        doc_gen = DocGenerator(pipeline)
        return doc_gen.generate_docstring(code_snippet)

    # Retrieve similar examples from the index
    results = index_manager.search(code_snippet, k=top_k)

    context_str = ""
    for i, ctx in enumerate(results[:top_k]):
        question = ctx.get('question', '')
        code = ctx.get('generated_mongodb_query', ctx.get('mongodb_query', ''))
        if isinstance(code, list) and len(code) > 0:
            code = code[0]
        context_str += f"--- Reference {i + 1} ---\n"
        if question:
            context_str += f"Description: {question}\n"
        if code:
            context_str += f"Code: {code}\n\n"

    # Build a documentation-specific prompt with retrieved context
    prompt = (
        "You are an expert programmer. Generate a concise, accurate, and professional "
        "docstring/description for the given code.\n\n"
        f"Reference Examples:\n{context_str}\n"
        f"Code:\n{code_snippet}\n\n"
        "Description:"
    )

    return pipeline.generate(prompt)


# ==========================================
# 5. 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()

    # ---- TEXT-TO-MONGODB ----
    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\nTask: Text-to-MongoDB\nModel: {model_choice}\nRAG: Enabled\n\n{context_str}"
        else:
            # Route to correct model based on selection
            if model_choice == "Fine-Tuned (LoRA)":
                result = finetuned_mongo_generator.generate_single(input_text)
            else:
                result = base_mongo_generator.generate_single(input_text)

            latency = result['latency_ms']
            info_text = f"⏱️ Latency: {latency:.2f} ms\nTask: Text-to-MongoDB\nModel: {model_choice}\nRAG: Disabled"
            return result['generated_query'], info_text

    # ---- DOCUMENTATION GENERATION ----
    elif task == "Documentation Generation":
        if use_rag and index_manager is not None:
            # Route to correct model + RAG
            if model_choice == "Fine-Tuned (LoRA)":
                generated_doc = generate_doc_with_rag(input_text, finetuned_pipeline, top_k=int(top_k))
            else:
                generated_doc = generate_doc_with_rag(input_text, base_pipeline, top_k=int(top_k))
            rag_status = "Enabled"
        else:
            # Route to correct model without RAG
            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)
            rag_status = "Disabled"

        latency = (time.time() - start_time) * 1000
        info_text = (
            f"⏱️ Latency: {latency:.2f} ms\n"
            f"Task: Documentation Generation\n"
            f"Model: {model_choice}\n"
            f"RAG: {rag_status}"
        )
        return generated_doc, info_text


# ==========================================
# 6. Dynamic UI Update Function
# ==========================================
def update_ui(task):
    """Dynamically updates input, output, AND RAG checkbox label based on 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"),
            gr.update(label="🧠 Use RAG (Text-to-MongoDB only)")
        )
    else:
        return (
            gr.update(
                placeholder="e.g., def add(a, b):\n    return a + b",
                label="Code Snippet"
            ),
            gr.update(language="python"),
            gr.update(label="🧠 Use RAG (Documentation only)")
        )


# ==========================================
# 7. 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, use_rag]
    )

    # Generate button click
    generate_btn.click(
        fn=generate_output,
        inputs=[task_selector, input_box, model_choice, use_rag, top_k_slider],
        outputs=[output_box, info_output]
    )

# ==========================================
# 8. Launch the App
# ==========================================
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )