Update app.py

app.py

@@ -414,136 +414,135 @@ class Phi2EducationalLLM(Runnable):
             return f"[Error generating response: {str(e)}]"
 
     def stream_generate(self, input: Input, config=None):
-    """Streaming generation method for real-time response display."""
-    import time
-    from datetime import datetime
+        """Streaming generation method for real-time response display."""
+        import time
+        from datetime import datetime
 
-    start_stream_time = time.perf_counter()
-    current_time = datetime.now()
+        start_stream_time = time.perf_counter()
+        current_time = datetime.now()
 
-    # --- Debug Start ---
-    logger.info("Starting stream_generate...")
-    logger.debug(f"Input received: {input}")
-    # -------------------
+        # --- Debug Start ---
+        logger.info("Starting stream_generate...")
+        logger.debug(f"Input received: {input}")
+        # -------------------
 
-    # Handle input
-    if isinstance(input, dict):
-        prompt = input.get('input', str(input))
-    else:
-        prompt = str(input)
+        # Handle input
+        if isinstance(input, dict):
+            prompt = input.get('input', str(input))
+        else:
+            prompt = str(input)
 
-    try:
-        # === Configurable Generation Parameters ===
-        temperature = config.get("temperature", 0.7) if config else 0.7
-        top_k = config.get("top_k", 50) if config else 50
-        top_p = config.get("top_p", 0.9) if config else 0.9
-        max_new_tokens = config.get("max_new_tokens", 600) if config else 600
-        timeout_seconds = config.get("timeout_seconds", 15) if config else 15
-
-        # === Prompt Construction ===
         try:
-            messages = [
-                {"role": "system", "content": SYSTEM_PROMPT},
-                {"role": "user", "content": prompt}
-            ]
-            text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        except Exception as e:
-            logger.warning(f"Failed to use chat template: {e}")
-            if "phi" in self.model_name.lower():
-                text = f"Instruct: {SYSTEM_PROMPT}\n\nUser: {prompt}\nOutput:"
-            else:
-                text = f"<|system|>\n{SYSTEM_PROMPT}<|end|>\n<|user|>\n{prompt}<|end|>\n<|assistant|>\n"
-
-        # === Tokenize ===
-        inputs = self.tokenizer(
-            [text], 
-            return_tensors="pt", 
-            padding=True, 
-            truncation=True, 
-            max_length=self.tokenizer.model_max_length
-        )
-        if torch.cuda.is_available():
-            inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
-
-        input_ids = inputs["input_ids"]
-        attention_mask = inputs["attention_mask"]
-        input_length = input_ids.shape[1]
-
-        # === Streaming Generation ===
-        generated_tokens = []
-        past_key_values = None
-        eos_token_id = self.tokenizer.eos_token_id
-        start_time = time.time()
-
-        logger.info("Beginning token-by-token generation...")
-
-        for step in range(max_new_tokens):
-            if time.time() - start_time > timeout_seconds:
-                logger.warning("Timeout reached. Ending stream.")
-                break
-
-            with torch.no_grad():
-                model_inputs = {
-                    "attention_mask": attention_mask,
-                    "use_cache": True
-                }
-
-                if past_key_values is None:
-                    model_inputs["input_ids"] = input_ids
+            # === Configurable Generation Parameters ===
+            temperature = config.get("temperature", 0.7) if config else 0.7
+            top_k = config.get("top_k", 50) if config else 50
+            top_p = config.get("top_p", 0.9) if config else 0.9
+            max_new_tokens = config.get("max_new_tokens", 600) if config else 600
+            timeout_seconds = config.get("timeout_seconds", 15) if config else 15
+
+            # === Prompt Construction ===
+            try:
+                messages = [
+                    {"role": "system", "content": SYSTEM_PROMPT},
+                    {"role": "user", "content": prompt}
+                ]
+                text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            except Exception as e:
+                logger.warning(f"Failed to use chat template: {e}")
+                if "phi" in self.model_name.lower():
+                    text = f"Instruct: {SYSTEM_PROMPT}\n\nUser: {prompt}\nOutput:"
                 else:
-                    model_inputs["input_ids"] = next_token
-                    model_inputs["past_key_values"] = past_key_values
-
-                outputs = self.model(**model_inputs)
-                logits = outputs.logits[:, -1, :]
-                past_key_values = outputs.past_key_values
-
-                # Sampling
-                logits = logits / temperature
-                filtered_logits = self._top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
-                probs = torch.nn.functional.softmax(filtered_logits, dim=-1)
-                next_token = torch.multinomial(probs, num_samples=1)
+                    text = f"<|system|>\n{SYSTEM_PROMPT}<|end|>\n<|user|>\n{prompt}<|end|>\n<|assistant|>\n"
 
-                token_id = next_token.item()
-                logger.debug(f"Step {step}: Token ID = {token_id}")
+            # === Tokenize ===
+            inputs = self.tokenizer(
+                [text], 
+                return_tensors="pt", 
+                padding=True, 
+                truncation=True, 
+                max_length=self.tokenizer.model_max_length
+            )
+            if torch.cuda.is_available():
+                inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
 
-                if eos_token_id is not None and token_id == eos_token_id:
-                    logger.info("EOS token encountered. Ending generation.")
-                    break
+            input_ids = inputs["input_ids"]
+            attention_mask = inputs["attention_mask"]
+            input_length = input_ids.shape[1]
 
-                generated_tokens.append(token_id)
+            # === Streaming Generation ===
+            generated_tokens = []
+            past_key_values = None
+            eos_token_id = self.tokenizer.eos_token_id
+            start_time = time.time()
 
-                # Decode efficiently
-                new_text = self.tokenizer.decode([token_id], skip_special_tokens=True)
-                yield new_text
+            logger.info("Beginning token-by-token generation...")
 
-                # Optional heuristic: stop on sentence-ending punctuation
-                if new_text.strip().endswith(('.', '?', '!')):
-                    logger.info("Sentence-ending punctuation hit. Ending early.")
+            for step in range(max_new_tokens):
+                if time.time() - start_time > timeout_seconds:
+                    logger.warning("Timeout reached. Ending stream.")
                     break
 
-                # Prepare for next step
-                input_ids = next_token
-                attention_mask = torch.cat([
-                    attention_mask,
-                    torch.ones((1, 1), dtype=attention_mask.dtype, device=attention_mask.device)
-                ], dim=-1)
-
-        # Final output logging
-        final_text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
-        logger.info(f"Streaming complete. Tokens generated: {len(generated_tokens)}")
-
-        end_stream_time = time.perf_counter()
-        stream_time = end_stream_time - start_stream_time
-        log_metric(f"LLM Stream time: {stream_time:0.4f} seconds. Tokens generated: {len(generated_tokens)}. Model: {self.model_name}. Timestamp: {current_time:%Y-%m-%d %H:%M:%S}")
-
-    except Exception as e:
-        logger.error(f"Streaming generation error: {e}")
-        end_stream_time = time.perf_counter()
-        stream_time = end_stream_time - start_stream_time
-        log_metric(f"LLM Stream time (error): {stream_time:0.4f} seconds. Model: {self.model_name}. Timestamp: {current_time:%Y-%m-%d %H:%M:%S}")
-        yield f"[Error in streaming generation: {str(e)}]"
+                with torch.no_grad():
+                    model_inputs = {
+                        "attention_mask": attention_mask,
+                        "use_cache": True
+                    }
+
+                    if past_key_values is None:
+                        model_inputs["input_ids"] = input_ids
+                    else:
+                        model_inputs["input_ids"] = next_token
+                        model_inputs["past_key_values"] = past_key_values
+
+                    outputs = self.model(**model_inputs)
+                    logits = outputs.logits[:, -1, :]
+                    past_key_values = outputs.past_key_values
+
+                    # Sampling
+                    logits = logits / temperature
+                    filtered_logits = self._top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
+                    probs = torch.nn.functional.softmax(filtered_logits, dim=-1)
+                    next_token = torch.multinomial(probs, num_samples=1)
+
+                    token_id = next_token.item()
+                    logger.debug(f"Step {step}: Token ID = {token_id}")
+
+                    if eos_token_id is not None and token_id == eos_token_id:
+                        logger.info("EOS token encountered. Ending generation.")
+                        break
+
+                    generated_tokens.append(token_id)
+
+                    # Decode efficiently
+                    new_text = self.tokenizer.decode([token_id], skip_special_tokens=True)
+                    yield new_text
+
+                    # Optional heuristic: stop on sentence-ending punctuation
+                    if new_text.strip().endswith(('.', '?', '!')):
+                        logger.info("Sentence-ending punctuation hit. Ending early.")
+                        break
+
+                    # Prepare for next step
+                    input_ids = next_token
+                    attention_mask = torch.cat([
+                        attention_mask,
+                        torch.ones((1, 1), dtype=attention_mask.dtype, device=attention_mask.device)
+                    ], dim=-1)
+
+            # Final output logging
+            final_text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
+            logger.info(f"Streaming complete. Tokens generated: {len(generated_tokens)}")
+
+            end_stream_time = time.perf_counter()
+            stream_time = end_stream_time - start_stream_time
+            log_metric(f"LLM Stream time: {stream_time:0.4f} seconds. Tokens generated: {len(generated_tokens)}. Model: {self.model_name}. Timestamp: {current_time:%Y-%m-%d %H:%M:%S}")
 
+        except Exception as e:
+            logger.error(f"Streaming generation error: {e}")
+            end_stream_time = time.perf_counter()
+            stream_time = end_stream_time - start_stream_time
+            log_metric(f"LLM Stream time (error): {stream_time:0.4f} seconds. Model: {self.model_name}. Timestamp: {current_time:%Y-%m-%d %H:%M:%S}")
+            yield f"[Error in streaming generation: {str(e)}]"
 
     def _top_k_top_p_filtering(self, logits, top_k=50, top_p=0.9):
         """Apply top-k and top-p filtering to logits"""
@@ -880,7 +879,7 @@ mathjax_config = '''
 window.MathJax = {
   tex: {
     inlineMath: [['\\\\(', '\\\\)']],
-    displayMath: [[', '], ['\\\\[', '\\\\]']],
+    displayMath: [['$', '$'], ['\\\\[', '\\\\]']],
     packages: {'[+]': ['ams']}
   },
   svg: {fontCache: 'global'},
@@ -1078,7 +1077,7 @@ def create_interface():
         
         with gr.Column(elem_classes=["main-container"]):
             # Title Section
-            gr.HTML('<div class="title-header"><h1> Mimir 🎓</h1></div>')
+            gr.HTML('<div class="title-header"><h1>🎓 Mimir 🎓</h1></div>')
             
             # Chat Section
             with gr.Row():