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 )