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):
        try:
            # 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,
                trust_remote_code=True
            )
            self.code_model = AutoModelForCausalLM.from_pretrained(
                self.code_model_name,
                torch_dtype=torch.bfloat16,
                device_map="auto",
                trust_remote_code=True
            )

            # Initialize Docs Assistant (Using Zephyr instead)
            print("Loading Documentation Assistant Model...")
            self.docs_model_name = "HuggingFaceH4/zephyr-7b-beta"
            self.docs_tokenizer = AutoTokenizer.from_pretrained(
                self.docs_model_name,
                trust_remote_code=True
            )
            self.docs_model = AutoModelForCausalLM.from_pretrained(
                self.docs_model_name,
                torch_dtype=torch.bfloat16,
                device_map="auto",
                trust_remote_code=True
            )

            # 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,
                trust_remote_code=True
            )
            self.pdf_model = AutoModelForCausalLM.from_pretrained(
                self.pdf_model_name,
                torch_dtype=torch.bfloat16,
                device_map="auto",
                trust_remote_code=True
            )
            
            print("All models loaded successfully!")
            
        except Exception as e:
            print(f"Error initializing models: {str(e)}")
            raise RuntimeError(f"Failed to initialize one or more models: {str(e)}")
    
    @spaces.GPU
    def process_code_query(self, query):
        try:
            if not query.strip():
                return "Please enter a coding question."
                
            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:
            print(f"Code query error: {str(e)}")
            return f"Error processing code query: {str(e)}"

    @spaces.GPU
    def process_docs_query(self, query, doc_file):
        try:
            if not query.strip():
                return "Please enter a documentation query."
            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:
            print(f"Documentation query error: {str(e)}")
            return f"Error processing documentation query: {str(e)}"

    @spaces.GPU
    def process_pdf_query(self, query, pdf_file):
        try:
            if not query.strip():
                return "Please enter a question about the PDF."
            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:
            print(f"PDF query error: {str(e)}")
            return f"Error processing PDF query: {str(e)}"

    def _read_file_content(self, file):
        try:
            content = ""
            if file.name.endswith('.pdf'):
                content = self._extract_pdf_text(file)
            else:
                content = file.read().decode('utf-8')
            return content
        except Exception as e:
            print(f"File reading error: {str(e)}")
            raise

    def _extract_pdf_text(self, pdf_file):
        try:
            reader = PdfReader(pdf_file)
            text = ""
            for page in reader.pages:
                text += page.extract_text() + "\n"
            return text
        except Exception as e:
            print(f"PDF extraction error: {str(e)}")
            raise

# 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()