Update app.py

app.py

@@ -18,6 +18,8 @@ from huggingface_hub import HfApi, HfFolder
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import time
+import logging
+import asyncio
 
 # ========== CONFIGURATION ==========
 PROFILES_DIR = "student_profiles"
@@ -28,6 +30,9 @@ MAX_AGE = 120
 SESSION_TOKEN_LENGTH = 32
 HF_TOKEN = os.getenv("HF_TOKEN")
 
+# Initialize logging
+logging.basicConfig(filename='app.log', level=logging.INFO)
+
 # Model configuration
 MODEL_CHOICES = {
     "TinyLlama (Fastest)": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
@@ -59,38 +64,56 @@ class ModelLoader:
         self.loading = True
         self.error = None
         try:
-            progress(0, desc=f"Loading {model_name}...")
+            progress(0.1, desc="Initializing...")
             
             # Clear previous model if any
             if self.model:
                 del self.model
                 del self.tokenizer
                 torch.cuda.empty_cache()
+                time.sleep(2)  # Allow CUDA cleanup
+            
+            # Load with optimized settings
+            model_kwargs = {
+                "trust_remote_code": True,
+                "torch_dtype": torch.float16,
+                "device_map": "auto",
+                "low_cpu_mem_usage": True
+            }
+            
+            if "TinyLlama" in model_name:
+                model_kwargs["attn_implementation"] = "flash_attention_2"
             
-            # Load tokenizer first
+            progress(0.3, desc="Loading tokenizer...")
             self.tokenizer = AutoTokenizer.from_pretrained(
                 MODEL_CHOICES[model_name],
                 trust_remote_code=True
             )
-            progress(0.3, desc="Loaded tokenizer...")
             
-            # Load model with appropriate settings
+            progress(0.6, desc="Loading model...")
             self.model = AutoModelForCausalLM.from_pretrained(
                 MODEL_CHOICES[model_name],
-                trust_remote_code=True,
-                torch_dtype=torch.float16,
-                device_map="auto" if torch.cuda.is_available() else None,
-                low_cpu_mem_usage=True
+                **model_kwargs
             )
             
+            # Verify model responsiveness
+            progress(0.8, desc="Verifying model...")
+            test_input = self.tokenizer("Test", return_tensors="pt").to(self.model.device)
+            _ = self.model.generate(**test_input, max_new_tokens=1)
+            
+            self.model.eval()  # Disable dropout
             progress(0.9, desc="Finalizing...")
             self.loaded = True
             self.current_model = model_name
             return self.model, self.tokenizer
             
+        except torch.cuda.OutOfMemoryError:
+            self.error = "Out of GPU memory. Try a smaller model."
+            logging.error(self.error)
+            return None, None
         except Exception as e:
             self.error = str(e)
-            print(f"Error loading model: {self.error}")
+            logging.error(f"Model loading error: {self.error}")
             return None, None
         finally:
             self.loading = False
@@ -132,7 +155,7 @@ def validate_age(age: Union[int, float, str]) -> int:
 def validate_file(file_obj) -> None:
     """Validate uploaded file."""
     if not file_obj:
-        raise gr.Error("No file uploaded")
+        raise ValueError("No file uploaded")
     
     file_ext = os.path.splitext(file_obj.name)[1].lower()
     if file_ext not in ALLOWED_FILE_TYPES:
@@ -157,7 +180,7 @@ def extract_text_from_file(file_path: str, file_ext: str) -> str:
                 if not text.strip():
                     raise ValueError("PyMuPDF returned empty text")
             except Exception as e:
-                print(f"PyMuPDF failed, trying OCR fallback: {str(e)}")
+                logging.warning(f"PyMuPDF failed: {str(e)}. Trying OCR fallback...")
                 text = extract_text_from_pdf_with_ocr(file_path)
         
         elif file_ext in ['.png', '.jpg', '.jpeg']:
@@ -172,7 +195,8 @@ def extract_text_from_file(file_path: str, file_ext: str) -> str:
         return text
     
     except Exception as e:
-        raise gr.Error(f"Text extraction error: {str(e)}")
+        logging.error(f"Text extraction error: {str(e)}")
+        raise gr.Error(f"Text extraction error: {str(e)}\nTips: Use high-quality images/PDFs with clear text.")
 
 def extract_text_from_pdf_with_ocr(file_path: str) -> str:
     """Fallback PDF text extraction using OCR."""
@@ -182,7 +206,10 @@ def extract_text_from_pdf_with_ocr(file_path: str) -> str:
         for page in doc:
             pix = page.get_pixmap()
             img = Image.open(io.BytesIO(pix.tobytes()))
-            text += pytesseract.image_to_string(img) + '\n'
+            # Preprocess image for better OCR
+            img = img.convert('L')  # Grayscale
+            img = img.point(lambda x: 0 if x < 128 else 255)  # Binarize
+            text += pytesseract.image_to_string(img, config='--psm 6 --oem 3') + '\n'
     except Exception as e:
         raise ValueError(f"PDF OCR failed: {str(e)}")
     return text
@@ -192,7 +219,7 @@ def extract_text_with_ocr(file_path: str) -> str:
     try:
         image = Image.open(file_path)
         
-        # Preprocess image for better OCR results
+        # Enhanced preprocessing
         image = image.convert('L')  # Convert to grayscale
         image = image.point(lambda x: 0 if x < 128 else 255, '1')  # Thresholding
         
@@ -355,6 +382,10 @@ class TranscriptParser:
             "completion_status": self._calculate_completion()
         }, indent=2)
 
+async def parse_transcript_async(file_obj, progress=gr.Progress()):
+    """Async wrapper for transcript parsing"""
+    return await asyncio.to_thread(parse_transcript, file_obj, progress)
+
 def parse_transcript_with_ai(text: str, progress=gr.Progress()) -> Dict:
     """Use AI model to parse transcript text with progress feedback"""
     model, tokenizer = model_loader.load_model(model_loader.current_model or DEFAULT_MODEL, progress)
@@ -393,7 +424,7 @@ def parse_transcript_with_ai(text: str, progress=gr.Progress()) -> Dict:
         return validate_parsed_data(formatted_data)
     
     except Exception as e:
-        print(f"Structured parsing failed, falling back to AI: {str(e)}")
+        logging.warning(f"Structured parsing failed, falling back to AI: {str(e)}")
         # Fall back to AI parsing if structured parsing fails
         return parse_transcript_with_ai_fallback(text, progress)
 
@@ -424,10 +455,10 @@ def parse_transcript_with_ai_fallback(text: str, progress=gr.Progress()) -> Dict
         progress(0.1, desc="Processing transcript with AI...")
         
         # Tokenize and generate response
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        inputs = model_loader.tokenizer(prompt, return_tensors="pt").to(model_loader.model.device)
         progress(0.4)
         
-        outputs = model.generate(
+        outputs = model_loader.model.generate(
             **inputs,
             max_new_tokens=1500,
             temperature=0.1,
@@ -436,20 +467,24 @@ def parse_transcript_with_ai_fallback(text: str, progress=gr.Progress()) -> Dict
         progress(0.8)
         
         # Decode the response
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        response = model_loader.tokenizer.decode(outputs[0], skip_special_tokens=True)
         
         # Extract JSON from response
-        json_str = response.split('```json')[1].split('```')[0].strip() if '```json' in response else response
+        try:
+            json_str = response.split('```json')[1].split('```')[0].strip()
+            parsed_data = json.loads(json_str)
+        except (IndexError, json.JSONDecodeError):
+            # Fallback: Extract JSON-like substring
+            json_str = re.search(r'\{.*\}', response, re.DOTALL).group()
+            parsed_data = json.loads(json_str)
         
-        # Parse and validate
-        parsed_data = json.loads(json_str)
         progress(1.0)
-        
         return validate_parsed_data(parsed_data)
     
     except torch.cuda.OutOfMemoryError:
         raise gr.Error("The model ran out of memory. Try with a smaller transcript or use a smaller model.")
     except Exception as e:
+        logging.error(f"AI parsing error: {str(e)}")
         raise gr.Error(f"Error processing transcript: {str(e)}")
 
 def validate_parsed_data(data: Dict) -> Dict:
@@ -546,6 +581,7 @@ def parse_transcript(file_obj, progress=gr.Progress()) -> Tuple[str, Optional[Di
         return output_text, transcript_data
     
     except Exception as e:
+        logging.error(f"Transcript processing error: {str(e)}")
         return f"Error processing transcript: {str(e)}", None
 
 # ========== LEARNING STYLE QUIZ ==========
@@ -801,11 +837,12 @@ class ProfileManager:
                         repo_type="dataset"
                     )
                 except Exception as e:
-                    print(f"Failed to upload to HF Hub: {str(e)}")
+                    logging.error(f"Failed to upload to HF Hub: {str(e)}")
             
             return self._generate_profile_summary(data)
         
         except Exception as e:
+            logging.error(f"Error saving profile: {str(e)}")
             raise gr.Error(f"Error saving profile: {str(e)}")
     
     def load_profile(self, name: str = None, session_token: str = None) -> Dict:
@@ -850,7 +887,7 @@ class ProfileManager:
                 return json.load(f)
         
         except Exception as e:
-            print(f"Error loading profile: {str(e)}")
+            logging.error(f"Error loading profile: {str(e)}")
             return {}
     
     def list_profiles(self, session_token: str = None) -> List[str]:
@@ -879,7 +916,7 @@ class ProfileManager:
         markdown = f"""## Student Profile: {data['name']}
 ### Basic Information
 - **Age:** {data['age']}
-- **Interests:** {data['interests']}
+- **Interests:** {data.get('interests', 'Not specified')}
 - **Learning Style:** {learning_style.split('##')[0].strip()}
 ### Academic Information
 {self._format_transcript(transcript)}
@@ -965,7 +1002,7 @@ class TeachingAssistant:
             return response
         
         except Exception as e:
-            print(f"Error generating response: {str(e)}")
+            logging.error(f"Error generating response: {str(e)}")
             return "I encountered an error processing your request. Please try again."
     
     def _update_context(self, message: str, history: List[List[Union[str, None]]]) -> None:
@@ -1327,21 +1364,33 @@ def create_interface():
                         transcript_data = gr.State()
                 
                 def process_transcript_and_update(file_obj, current_tab_status, progress=gr.Progress()):
-                    output_text, data = parse_transcript(file_obj, progress)
-                    if "Error" not in output_text:
-                        new_status = current_tab_status.copy()
-                        new_status[0] = True
-                        return output_text, data, new_status, \
-                               gr.update(elem_classes="completed-tab"), \
-                               gr.update(interactive=True), \
-                               gr.update(visible=False)
-                    return output_text, data, current_tab_status, \
-                           gr.update(), gr.update(), gr.update()
+                    try:
+                        output_text, data = parse_transcript(file_obj, progress)
+                        if "Error" not in output_text:
+                            new_status = current_tab_status.copy()
+                            new_status[0] = True
+                            return (
+                                output_text, 
+                                data, 
+                                new_status,
+                                gr.update(elem_classes="completed-tab"),
+                                gr.update(interactive=True),
+                                gr.update(visible=False)
+                    except Exception as e:
+                        logging.error(f"Upload error: {str(e)}")
+                    return (
+                        "Error processing transcript. Please try again.", 
+                        None, 
+                        current_tab_status,
+                        gr.update(),
+                        gr.update(),
+                        gr.update(visible=True, value=f"<div class='nav-message'>Error: {str(e)}</div>"))
                 
                 upload_btn.click(
                     fn=process_transcript_and_update,
                     inputs=[transcript_file, tab_completed],
-                    outputs=[transcript_output, transcript_data, tab_completed, step1, step2, nav_message]
+                    outputs=[transcript_output, transcript_data, tab_completed, step1, step2, nav_message],
+                    concurrency_limit=1
                 )
             
             # ===== TAB 2: Learning Style Quiz =====
@@ -1388,18 +1437,28 @@ def create_interface():
                     current_tab_status = args[0]
                     answers = args[1:]
                     
-                    result = learning_style_quiz.evaluate_quiz(*answers)
-                    
-                    new_status = current_tab_status.copy()
-                    new_status[1] = True
-                    
-                    return result, \
-                           gr.update(visible=True), \
-                           new_status, \
-                           gr.update(elem_classes="completed-tab"), \
-                           gr.update(interactive=True), \
-                           gr.update(value="<div class='alert-box'>Quiz submitted successfully! Scroll down to view your results.</div>", visible=True), \
-                           gr.update(visible=False)
+                    try:
+                        result = learning_style_quiz.evaluate_quiz(*answers)
+                        new_status = current_tab_status.copy()
+                        new_status[1] = True
+                        return (
+                            result,
+                            gr.update(visible=True),
+                            new_status,
+                            gr.update(elem_classes="completed-tab"),
+                            gr.update(interactive=True),
+                            gr.update(value="<div class='alert-box'>Quiz submitted successfully! Scroll down to view your results.</div>", visible=True),
+                            gr.update(visible=False))
+                    except Exception as e:
+                        logging.error(f"Quiz error: {str(e)}")
+                        return (
+                            f"Error evaluating quiz: {str(e)}",
+                            gr.update(visible=True),
+                            current_tab_status,
+                            gr.update(),
+                            gr.update(),
+                            gr.update(value=f"<div class='nav-message'>Error: {str(e)}</div>", visible=True),
+                            gr.update(visible=False))
                 
                 quiz_submit.click(
                     fn=submit_quiz_and_update,
@@ -1456,11 +1515,26 @@ def create_interface():
                         )
                 
                 def save_personal_info(name, age, interests, current_tab_status):
-                    if name.strip() and age and interests.strip():
+                    try:
+                        name = validate_name(name)
+                        age = validate_age(age)
+                        interests = sanitize_input(interests)
+                        
                         new_status = current_tab_status.copy()
                         new_status[2] = True
-                        return new_status, gr.update(elem_classes="completed-tab"), gr.update(interactive=True), gr.update(value="<div class='alert-box'>Information saved!</div>", visible=True), gr.update(visible=False)
-                    return current_tab_status, gr.update(), gr.update(), gr.update(visible=False), gr.update(visible=True)
+                        return (
+                            new_status,
+                            gr.update(elem_classes="completed-tab"),
+                            gr.update(interactive=True),
+                            gr.update(value="<div class='alert-box'>Information saved!</div>", visible=True),
+                            gr.update(visible=False))
+                    except Exception as e:
+                        return (
+                            current_tab_status,
+                            gr.update(),
+                            gr.update(),
+                            gr.update(visible=False),
+                            gr.update(visible=True, value=f"<div class='nav-message'>Error: {str(e)}</div>"))
                 
                 save_personal_btn.click(
                     fn=save_personal_info,
@@ -1502,14 +1576,28 @@ def create_interface():
                     inputs = args[:-1]  # All except the last which is tab_completed
                     current_tab_status = args[-1]
                     
-                    # Call the original save function
-                    summary = profile_manager.save_profile(*inputs)
-                    
-                    # Update completion status
-                    new_status = current_tab_status.copy()
-                    new_status[3] = True
-                    
-                    return summary, new_status, gr.update(elem_classes="completed-tab"), gr.update(interactive=True), gr.update(visible=False)
+                    try:
+                        # Call the original save function
+                        summary = profile_manager.save_profile(*inputs)
+                        
+                        # Update completion status
+                        new_status = current_tab_status.copy()
+                        new_status[3] = True
+                        
+                        return (
+                            summary,
+                            new_status,
+                            gr.update(elem_classes="completed-tab"),
+                            gr.update(interactive=True),
+                            gr.update(visible=False))
+                    except Exception as e:
+                        logging.error(f"Save profile error: {str(e)}")
+                        return (
+                            f"Error saving profile: {str(e)}",
+                            current_tab_status,
+                            gr.update(),
+                            gr.update(),
+                            gr.update(visible=True, value=f"<div class='nav-message'>Error: {str(e)}</div>"))
                 
                 save_btn.click(
                     fn=save_profile_and_update,
@@ -1524,10 +1612,10 @@ def create_interface():
                     fn=lambda: profile_manager.list_profiles(session_token.value),
                     outputs=load_profile_dropdown
                 ).then(
-                    fn=lambda: gr.update(visible=True),
+                    fn=lambda: gr.update(visible=bool(profile_manager.list_profiles(session_token.value))),
                     outputs=load_btn
                 ).then(
-                    fn=lambda: gr.update(visible=True),
+                    fn=lambda: gr.update(visible=bool(profile_manager.list_profiles(session_token.value))),
                     outputs=delete_btn
                 )
                 
@@ -1548,6 +1636,7 @@ def create_interface():
                             profile_path.unlink()
                         return "Profile deleted successfully", ""
                     except Exception as e:
+                        logging.error(f"Delete profile error: {str(e)}")
                         raise gr.Error(f"Error deleting profile: {str(e)}")
                 
                 delete_btn.click(
@@ -1608,41 +1697,44 @@ def create_interface():
         
         # Tab navigation logic with completion check
         def navigate_to_tab(tab_index: int, tab_completed_status):
-            # Always allow going back to previous tabs
             current_tab = tabs.selected
-            if current_tab is not None and tab_index > current_tab:
-                # Check if current tab is completed
-                if not tab_completed_status.get(current_tab, False):
-                    return gr.Tabs(selected=current_tab), \
-                           gr.update(value=f"<div class='nav-message'>Please complete the current tab before proceeding to tab {tab_index + 1}</div>", visible=True), \
-                           gr.update(visible=False)
             
-            return gr.Tabs(selected=tab_index), gr.update(visible=False), gr.update(visible=False)
+            # Allow backward navigation
+            if tab_index <= current_tab:
+                return gr.Tabs(selected=tab_index), gr.update(visible=False)
+            
+            # Check if current tab is completed
+            if not tab_completed_status.get(current_tab, False):
+                return (
+                    gr.Tabs(selected=current_tab),
+                    gr.update(value=f"⚠️ Complete Step {current_tab+1} first!", visible=True))
+            
+            return gr.Tabs(selected=tab_index), gr.update(visible=False)
         
         step1.click(
             fn=lambda idx, status: navigate_to_tab(idx, status),
             inputs=[gr.State(0), tab_completed],
-            outputs=[tabs, nav_message, quiz_alert]
+            outputs=[tabs, nav_message]
         )
         step2.click(
             fn=lambda idx, status: navigate_to_tab(idx, status),
             inputs=[gr.State(1), tab_completed],
-            outputs=[tabs, nav_message, quiz_alert]
+            outputs=[tabs, nav_message]
         )
         step3.click(
             fn=lambda idx, status: navigate_to_tab(idx, status),
             inputs=[gr.State(2), tab_completed],
-            outputs=[tabs, nav_message, quiz_alert]
+            outputs=[tabs, nav_message]
         )
         step4.click(
             fn=lambda idx, status: navigate_to_tab(idx, status),
             inputs=[gr.State(3), tab_completed],
-            outputs=[tabs, nav_message, quiz_alert]
+            outputs=[tabs, nav_message]
         )
         step5.click(
             fn=lambda idx, status: navigate_to_tab(idx, status),
             inputs=[gr.State(4), tab_completed],
-            outputs=[tabs, nav_message, quiz_alert]
+            outputs=[tabs, nav_message]
         )
         
         # Model loading functions
@@ -1654,6 +1746,7 @@ def create_interface():
                 else:
                     return gr.update(value=f"<div class='nav-message'>Failed to load model: {model_loader.error}</div>", visible=True)
             except Exception as e:
+                logging.error(f"Model loading error: {str(e)}")
                 return gr.update(value=f"<div class='nav-message'>Error: {str(e)}</div>", visible=True)
         
         load_model_btn.click(