app.py

import gradio as gr
import openai
from datasets import load_dataset
import logging
import time
from langchain.embeddings import HuggingFaceEmbeddings
import torch
import psutil
import GPUtil

# Set up logging with performance metrics
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def get_system_metrics():
    cpu_percent = psutil.cpu_percent()
    memory_percent = psutil.virtual_memory().percent
    if torch.cuda.is_available():
        gpu = GPUtil.getGPUs()[0]
        gpu_util = gpu.load * 100
        gpu_memory = gpu.memoryUtil * 100
    else:
        gpu_util = 0
        gpu_memory = 0
    return cpu_percent, memory_percent, gpu_util, gpu_memory

def process_query(query, dataset_choice="all"):
    start_time = time.time()
    try:
        # Original query processing code here...
        response = "Sample response"
        
        # Calculate performance metrics
        end_time = time.time()
        processing_time = end_time - start_time
        cpu_percent, memory_percent, gpu_util, gpu_memory = get_system_metrics()
        
        metrics = f"""
        Performance Metrics:
        Processing Time: {processing_time:.2f}s
        CPU Usage: {cpu_percent}%
        Memory Usage: {memory_percent}%
        GPU Utilization: {gpu_util:.1f}%
        GPU Memory: {gpu_memory:.1f}%
        """
        
        return response, metrics
        
    except Exception as e:
        return str(e), "Metrics unavailable"

# Enhanced Gradio interface with performance metrics
demo = gr.Interface(
    fn=process_query,
    inputs=[
        gr.Textbox(label="Question", placeholder="Ask any question..."),
        gr.Dropdown(
            choices=["all"] + dataset_names,
            label="Select Dataset",
            value="all"
        )
    ],
    outputs=[
        gr.Textbox(label="Response"),
        gr.Textbox(label="Performance Metrics")
    ],
    title="E5-Powered Multi-Dataset Knowledge Base",
    description="Search across RagBench datasets with real-time performance monitoring",
    analytics_enabled=True,
    examples=[
        ["What role does T-cell count play in severe human adenovirus type 55 (HAdV-55) infection?", "covidqa"],
        ["In what school district is Governor John R. Rogers High School located?", "hotpotqa"],
        ["What are the key financial metrics for Q3?", "finqa"]
    ]
)

if __name__ == "__main__":
    demo.queue()  # Enable queuing for performance monitoring
    demo.launch(debug=True, show_api=True)