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