Update app.py

app.py

@@ -10,39 +10,61 @@ 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"
-        )
+        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 (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 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,
+                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)
-        self.pdf_model = AutoModelForCausalLM.from_pretrained(
-            self.pdf_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,
+                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,
@@ -53,11 +75,14 @@ class UnifiedAssistant:
             )
             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."
             
@@ -73,11 +98,14 @@ class UnifiedAssistant:
             )
             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."
             
@@ -93,22 +121,31 @@ class UnifiedAssistant:
             )
             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):
-        content = ""
-        if file.name.endswith('.pdf'):
-            content = self._extract_pdf_text(file)
-        else:
-            content = file.read().decode('utf-8')
-        return content
+        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):
-        reader = PdfReader(pdf_file)
-        text = ""
-        for page in reader.pages:
-            text += page.extract_text() + "\n"
-        return text
+        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 = """