Spaces:

Emasters-CodeGen
/

CodeGen-Emasters-IIITH

Paused

Create app.py

86d6d13 verified 2 months ago

9.34 kB

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
    )

	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
	)