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 )