Spaces:

Rittttttu
/

CodeGen

Sleeping

Upload 42 files

7545355 verified 2 months ago

11.9 kB

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
    )

	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
	)