Update app.py

app.py

@@ -28,38 +28,75 @@ MAX_AGE = 120
 SESSION_TOKEN_LENGTH = 32
 HF_TOKEN = os.getenv("HF_TOKEN")
 
+# Model configuration
+MODEL_CHOICES = {
+    "TinyLlama (Fastest)": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+    "Phi-2 (Balanced)": "microsoft/phi-2",
+    "DeepSeek-V3 (Most Powerful)": "deepseek-ai/DeepSeek-V3"
+}
+DEFAULT_MODEL = "TinyLlama (Fastest)"
+
 # Initialize Hugging Face API
 if HF_TOKEN:
     hf_api = HfApi(token=HF_TOKEN)
     HfFolder.save_token(HF_TOKEN)
 
-# ========== DEEPSEEK MODEL LOADING ==========
-def load_deepseek_model():
-    """Load the DeepSeek model with progress tracking"""
-    try:
-        start_time = time.time()
-        tokenizer = AutoTokenizer.from_pretrained(
-            "deepseek-ai/DeepSeek-V3",
-            trust_remote_code=True
-        )
-        
-        model = AutoModelForCausalLM.from_pretrained(
-            "deepseek-ai/DeepSeek-V3",
-            trust_remote_code=True,
-            torch_dtype=torch.float16,
-            device_map="auto" if torch.cuda.is_available() else None
-        )
-        
-        load_time = time.time() - start_time
-        print(f"DeepSeek model loaded in {load_time:.2f} seconds")
-        return model, tokenizer
-        
-    except Exception as e:
-        print(f"Error loading DeepSeek model: {str(e)}")
-        return None, None
+# ========== OPTIMIZED MODEL LOADING ==========
+class ModelLoader:
+    def __init__(self):
+        self.model = None
+        self.tokenizer = None
+        self.loaded = False
+        self.loading = False
+        self.error = None
+        self.current_model = None
+    
+    def load_model(self, model_name, progress=gr.Progress()):
+        """Lazy load the model with progress feedback"""
+        if self.loaded and self.current_model == model_name:
+            return self.model, self.tokenizer
+            
+        self.loading = True
+        self.error = None
+        try:
+            progress(0, desc=f"Loading {model_name}...")
+            
+            # Clear previous model if any
+            if self.model:
+                del self.model
+                del self.tokenizer
+                torch.cuda.empty_cache()
+            
+            # Load tokenizer first
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                MODEL_CHOICES[model_name],
+                trust_remote_code=True
+            )
+            progress(0.3, desc="Loaded tokenizer...")
+            
+            # Load model with appropriate settings
+            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
+            )
+            
+            progress(0.9, desc="Finalizing...")
+            self.loaded = True
+            self.current_model = model_name
+            return self.model, self.tokenizer
+            
+        except Exception as e:
+            self.error = str(e)
+            print(f"Error loading model: {self.error}")
+            return None, None
+        finally:
+            self.loading = False
 
-# Load model at startup
-model, tokenizer = load_deepseek_model()
+# Initialize model loader
+model_loader = ModelLoader()
 
 # ========== UTILITY FUNCTIONS ==========
 def generate_session_token() -> str:
@@ -198,13 +235,14 @@ def remove_sensitive_info(text: str) -> str:
     return text
 
 # ========== TRANSCRIPT PARSING ==========
-def parse_transcript_with_deepseek(text: str) -> Dict:
-    """Use local DeepSeek model to parse transcript text"""
+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)
     if model is None or tokenizer is None:
-        raise gr.Error("DeepSeek model failed to load. Please try again later.")
+        raise gr.Error(f"Model failed to load. {model_loader.error or 'Please try loading a model first.'}")
     
     # Pre-process the text
-    text = remove_sensitive_info(text[:15000])  # Limit to first 15k chars
+    text = remove_sensitive_info(text[:15000])  # Limit input size
     
     prompt = f"""
     Analyze this academic transcript and extract structured information:
@@ -218,57 +256,41 @@ def parse_transcript_with_deepseek(text: str) -> Dict:
       * Credits earned
       * Year/semester taken
       * Grade level when taken
-    Return the data in this JSON structure:
-    {{
-        "grade_level": "11",
-        "gpa": {{
-            "weighted": "4.2",
-            "unweighted": "3.9"
-        }},
-        "courses": [
-            {{
-                "code": "MATH101",
-                "name": "Algebra II",
-                "grade": "A",
-                "credits": "1.0",
-                "year": "2023-2024",
-                "grade_level": "11"
-            }}
-        ]
-    }}
+    Return the data in JSON format.
+    
     Transcript Text:
     {text}
     """
     
     try:
+        progress(0.1, desc="Processing transcript...")
+        
         # Tokenize and generate response
         inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        progress(0.4)
         
         outputs = model.generate(
             **inputs,
-            max_new_tokens=2000,
+            max_new_tokens=1500,  # Reduced from original
             temperature=0.1,
             do_sample=True
         )
+        progress(0.8)
         
         # Decode the response
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        # Extract the JSON content from the response
-        if '```json' in response:
-            json_str = response.split('```json')[1].split('```')[0].strip()
-        elif '```' in response:
-            json_str = response.split('```')[1].split('```')[0].strip()
-        else:
-            json_str = response
+        # Extract JSON from response
+        json_str = response.split('```json')[1].split('```')[0].strip() if '```json' in response else response
         
-        # Parse and validate the JSON
+        # 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 upgrade your GPU.")
+        raise gr.Error("The model ran out of memory. Try with a smaller transcript or use a smaller model.")
     except Exception as e:
         raise gr.Error(f"Error processing transcript: {str(e)}")
 
@@ -333,7 +355,7 @@ def format_transcript_output(data: Dict) -> str:
     
     return '\n'.join(output)
 
-def parse_transcript(file_obj) -> Tuple[str, Optional[Dict]]:
+def parse_transcript(file_obj, progress=gr.Progress()) -> Tuple[str, Optional[Dict]]:
     """Main function to parse transcript files."""
     try:
         if not file_obj:
@@ -345,8 +367,8 @@ def parse_transcript(file_obj) -> Tuple[str, Optional[Dict]]:
         # Extract text from file
         text = extract_text_from_file(file_obj.name, file_ext)
         
-        # Use DeepSeek for parsing
-        parsed_data = parse_transcript_with_deepseek(text)
+        # Use AI for parsing
+        parsed_data = parse_transcript_with_ai(text, progress)
         
         # Format output text
         output_text = format_transcript_output(parsed_data)
@@ -1069,6 +1091,12 @@ def create_interface():
             background-color: #fff3e0;
             color: #e65100;
         }
+        .model-selection {
+            margin-bottom: 20px;
+            padding: 15px;
+            background: #f8f9fa;
+            border-radius: 8px;
+        }
         """
         
         gr.Markdown("""
@@ -1077,13 +1105,21 @@ def create_interface():
         Complete each step to get customized learning recommendations.
         """)
         
-        # Model loading status
-        model_status = gr.HTML(
-            value="<div class='model-loading'>Loading AI model... (This may take a few minutes)</div>" if model is None else "",
-            visible=model is None
-        )
-        
-        # Progress tracker - now with dynamic styling
+        # Model selection section
+        with gr.Group(elem_classes="model-selection"):
+            model_selector = gr.Dropdown(
+                choices=list(MODEL_CHOICES.keys()),
+                value=DEFAULT_MODEL,
+                label="Select AI Model",
+                interactive=True
+            )
+            load_model_btn = gr.Button("Load Selected Model", variant="secondary")
+            model_status = gr.HTML(
+                value="<div class='model-loading'>Model not loaded yet. Please select and load a model.</div>",
+                visible=True
+            )
+        
+        # Progress tracker
         with gr.Row():
             with gr.Column(scale=1):
                 step1 = gr.Button("1. Upload Transcript", elem_classes="incomplete-tab")
@@ -1132,11 +1168,8 @@ def create_interface():
                         )
                         transcript_data = gr.State()
                 
-                def process_transcript_and_update(file_obj, current_tab_status):
-                    if model is None:
-                        return "Error: AI model failed to load. Please try again later.", None, current_tab_status, gr.update(), gr.update(), gr.update()
-                    
-                    output_text, data = parse_transcript(file_obj)
+                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
@@ -1454,15 +1487,20 @@ def create_interface():
             outputs=[tabs, nav_message, quiz_alert]
         )
         
-        # Check model loading status
-        def check_model_status():
-            if model is not None and tokenizer is not None:
-                return gr.update(visible=False)
-            return gr.update(visible=True)
+        # Model loading functions
+        def load_selected_model(model_name, progress=gr.Progress()):
+            try:
+                model_loader.load_model(model_name, progress)
+                if model_loader.loaded:
+                    return gr.update(value=f"<div class='alert-box'>{model_name} loaded successfully!</div>", visible=True)
+                else:
+                    return gr.update(value=f"<div class='nav-message'>Failed to load model: {model_loader.error}</div>", visible=True)
+            except Exception as e:
+                return gr.update(value=f"<div class='nav-message'>Error: {str(e)}</div>", visible=True)
         
-        app.load(
-            fn=check_model_status,
-            inputs=None,
+        load_model_btn.click(
+            fn=load_selected_model,
+            inputs=model_selector,
             outputs=model_status
         )
     
@@ -1474,4 +1512,4 @@ app = create_interface()
 # For Hugging Face Spaces deployment
 if __name__ == "__main__":
     app.launch()
- 
+