Update app.py

app.py

@@ -1,87 +1,237 @@
-import gradio as gr
-from modules.code_assistant import CodeAssistant
-from modules.docs_assistant import DocsAssistant
-from modules.pdf_assistant import PDFAssistant
-from modules.utils import load_css
-
-def create_app():
-    # Initialize assistants
-    code_assistant = CodeAssistant()
-    docs_assistant = DocsAssistant()
-    pdf_assistant = PDFAssistant()
-    
-    with gr.Blocks(css=load_css()) as demo:
-        gr.Markdown("# Enterprise RAG Assistant")
-        
-        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")
-                    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")
-                    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")
-                    pdf_output = gr.Markdown()
-        
-        # Event handlers
-        code_submit.click(
-            code_assistant.generate_response,
-            inputs=[code_input],
-            outputs=[code_output]
-        )
-        
-        docs_submit.click(
-            docs_assistant.search_docs,
-            inputs=[docs_input, docs_file],
-            outputs=[docs_output]
-        )
-        
-        pdf_submit.click(
-            pdf_assistant.answer_query,
-            inputs=[pdf_query, pdf_file],
-            outputs=[pdf_output]
-        )
-    
-    return demo
-
-if __name__ == "__main__":
-    app = create_app()
+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()