app.py

import gradio as gr
import spaces
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from PyPDF2 import PdfReader

# Verify GPU availability
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

class UnifiedAssistant:
    def __init__(self):
        # Initialize Code Assistant (Qwen)
        print("Loading Code Assistant Model...")
        self.code_model_name = "Qwen/Qwen2.5-Coder-32B-Instruct"
        self.code_tokenizer = AutoTokenizer.from_pretrained(self.code_model_name)
        self.code_model = AutoModelForCausalLM.from_pretrained(
            self.code_model_name,
            torch_dtype=torch.bfloat16,
            device_map="auto"
        )

        # Initialize Docs Assistant (DocGPT)
        print("Loading Documentation Assistant Model...")
        self.docs_model_name = "Arc53/docsgpt-40b-falcon"
        self.docs_tokenizer = AutoTokenizer.from_pretrained(self.docs_model_name)
        self.docs_model = AutoModelForCausalLM.from_pretrained(
            self.docs_model_name,
            torch_dtype=torch.bfloat16,
            device_map="auto"
        )

        # Initialize PDF Assistant (Llama)
        print("Loading PDF Assistant Model...")
        self.pdf_model_name = "meta-llama/Llama-3.3-70B-Instruct"
        self.pdf_tokenizer = AutoTokenizer.from_pretrained(self.pdf_model_name)
        self.pdf_model = AutoModelForCausalLM.from_pretrained(
            self.pdf_model_name,
            torch_dtype=torch.bfloat16,
            device_map="auto"
        )
    
    @spaces.GPU
    def process_code_query(self, query):
        try:
            inputs = self.code_tokenizer(query, return_tensors="pt").to(self.code_model.device)
            outputs = self.code_model.generate(
                **inputs,
                max_length=2048,
                temperature=0.7,
                top_p=0.95,
                do_sample=True
            )
            return self.code_tokenizer.decode(outputs[0], skip_special_tokens=True)
        except Exception as e:
            return f"Error processing code query: {str(e)}"

    @spaces.GPU
    def process_docs_query(self, query, doc_file):
        try:
            if doc_file is None:
                return "Please upload a documentation file."
            
            doc_content = self._read_file_content(doc_file)
            prompt = f"Documentation: {doc_content}\nQuery: {query}"
            
            inputs = self.docs_tokenizer(prompt, return_tensors="pt").to(self.docs_model.device)
            outputs = self.docs_model.generate(
                **inputs,
                max_length=1024,
                temperature=0.3,
                top_p=0.95
            )
            return self.docs_tokenizer.decode(outputs[0], skip_special_tokens=True)
        except Exception as e:
            return f"Error processing documentation query: {str(e)}"

    @spaces.GPU
    def process_pdf_query(self, query, pdf_file):
        try:
            if pdf_file is None:
                return "Please upload a PDF file."
            
            pdf_text = self._extract_pdf_text(pdf_file)
            prompt = f"Context from PDF: {pdf_text}\nQuestion: {query}"
            
            inputs = self.pdf_tokenizer(prompt, return_tensors="pt").to(self.pdf_model.device)
            outputs = self.pdf_model.generate(
                **inputs,
                max_length=1024,
                temperature=0.3,
                top_p=0.95
            )
            return self.pdf_tokenizer.decode(outputs[0], skip_special_tokens=True)
        except Exception as e:
            return f"Error processing PDF query: {str(e)}"

    def _read_file_content(self, file):
        content = ""
        if file.name.endswith('.pdf'):
            content = self._extract_pdf_text(file)
        else:
            content = file.read().decode('utf-8')
        return content

    def _extract_pdf_text(self, pdf_file):
        reader = PdfReader(pdf_file)
        text = ""
        for page in reader.pages:
            text += page.extract_text() + "\n"
        return text

# Custom CSS for better UI
css = """
.gradio-container {
    font-family: 'Inter', sans-serif;
    max-width: 1200px !important;
    margin: auto;
}
.tabs {
    background: #f8f9fa;
    border-radius: 10px;
    padding: 20px;
    margin-bottom: 20px;
}
.input-box {
    border: 1px solid #e0e0e0;
    border-radius: 8px;
    padding: 12px;
}
.button {
    background: #2d63c8 !important;
    color: white !important;
    border-radius: 6px !important;
    padding: 10px 20px !important;
    transition: all 0.3s ease !important;
}
.button:hover {
    background: #1e4a9d !important;
    transform: translateY(-1px) !important;
}
.output-box {
    background: #ffffff;
    border: 1px solid #e0e0e0;
    border-radius: 8px;
    padding: 16px;
    margin-top: 12px;
}
"""

def create_app():
    print("Initializing RAG Assistant...")
    assistant = UnifiedAssistant()
    
    with gr.Blocks(css=css) as demo:
        gr.Markdown("""
        # 🚀 Enterprise RAG Assistant
        ### Your AI-Powered Documentation & Code Assistant
        
        This application combines three powerful AI models:
        - 💻 Code Assistant (Qwen2.5-Coder-32B)
        - 📚 Documentation Helper (DocGPT-40B)
        - 📑 PDF Analyzer (Llama-3.3-70B)
        """)
        
        with gr.Tabs() as tabs:
            # Code Assistant Tab
            with gr.Tab("💻 Code Assistant", id=1):
                with gr.Row():
                    with gr.Column():
                        code_input = gr.Textbox(
                            label="Ask coding questions",
                            placeholder="Enter your coding question...",
                            lines=3
                        )
                        code_submit = gr.Button("🔍 Get Code Solution", variant="primary")
                    code_output = gr.Code(
                        label="Code Output",
                        language="python"
                    )
            
            # Documentation Assistant Tab
            with gr.Tab("📚 Documentation Assistant", id=2):
                with gr.Row():
                    with gr.Column():
                        docs_input = gr.Textbox(
                            label="Documentation Query",
                            placeholder="Ask about technical documentation...",
                            lines=3
                        )
                        docs_file = gr.File(
                            label="Upload Documentation",
                            file_types=[".pdf", ".txt", ".md"]
                        )
                        docs_submit = gr.Button("🔍 Search Documentation", variant="primary")
                    docs_output = gr.Markdown()
            
            # PDF RAG Assistant Tab
            with gr.Tab("📑 PDF Assistant", id=3):
                with gr.Row():
                    with gr.Column():
                        pdf_file = gr.File(
                            label="Upload PDF",
                            file_types=[".pdf"]
                        )
                        pdf_query = gr.Textbox(
                            label="Ask about the PDF",
                            placeholder="Enter your question about the PDF...",
                            lines=3
                        )
                        pdf_submit = gr.Button("🔍 Get Answer", variant="primary")
                    pdf_output = gr.Markdown()

        # Event handlers
        code_submit.click(
            assistant.process_code_query,
            inputs=[code_input],
            outputs=[code_output]
        )
        
        docs_submit.click(
            assistant.process_docs_query,
            inputs=[docs_input, docs_file],
            outputs=[docs_output]
        )
        
        pdf_submit.click(
            assistant.process_pdf_query,
            inputs=[pdf_query, pdf_file],
            outputs=[pdf_output]
        )
    
    return demo

if __name__ == "__main__":
    app = create_app()
    app.launch()