Update app.py

app.py

@@ -414,95 +414,136 @@ 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"""
-        start_stream_time = time.perf_counter()
-        current_time = datetime.now()
-        
-        # Handle both string and dict inputs for flexibility
-        if isinstance(input, dict):
-            prompt = input.get('input', str(input))
-        else:
-            prompt = str(input)
-            
+    """Streaming generation method for real-time response display."""
+    import time
+    from datetime import datetime
+
+    start_stream_time = time.perf_counter()
+    current_time = datetime.now()
+
+    # --- 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)
+
+    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:
-            # Prepare input text
-            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:
-                if "phi" in self.model_name.lower():
-                    text = f"Instruct: {SYSTEM_PROMPT}\n\nUser: {prompt}\nOutput:"
+            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
                 else:
-                    text = f"<|system|>\n{SYSTEM_PROMPT}<|end|>\n<|user|>\n{prompt}<|end|>\n<|assistant|>\n"
+                    model_inputs["input_ids"] = next_token
+                    model_inputs["past_key_values"] = past_key_values
 
-            inputs = self.tokenizer([text], return_tensors="pt", padding=True, truncation=True, max_length=1024)
-            if torch.cuda.is_available():
-                inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
+                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)}]"
 
-            # Initialize for streaming
-            generated_tokens = []
-            input_length = inputs.input_ids.shape[1]
-            max_new_tokens = 600
-            
-            # Generate token by token
-            current_input_ids = inputs.input_ids
-            current_attention_mask = inputs.attention_mask
-            
-            for step in range(max_new_tokens):
-                with torch.no_grad():
-                    outputs = self.model(
-                        input_ids=current_input_ids,
-                        attention_mask=current_attention_mask,
-                        use_cache=True
-                    )
-                    
-                    # Get next token probabilities
-                    next_token_logits = outputs.logits[:, -1, :]
-                    
-                    # Apply temperature and sampling
-                    next_token_logits = next_token_logits / 0.7
-                    
-                    # Apply top-k and top-p filtering
-                    filtered_logits = self._top_k_top_p_filtering(next_token_logits, top_k=50, top_p=0.9)
-                    
-                    # Sample next token
-                    probs = torch.nn.functional.softmax(filtered_logits, dim=-1)
-                    next_token = torch.multinomial(probs, num_samples=1)
-                    
-                    # Check for end of sequence
-                    if next_token.item() == self.tokenizer.eos_token_id:
-                        break
-                    
-                    # Add to generated tokens
-                    generated_tokens.append(next_token.item())
-                    
-                    # Decode and yield partial result
-                    partial_text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
-                    yield partial_text
-                    
-                    # Update input for next iteration
-                    current_input_ids = torch.cat([current_input_ids, next_token], dim=-1)
-                    current_attention_mask = torch.cat([
-                        current_attention_mask, 
-                        torch.ones((1, 1), dtype=current_attention_mask.dtype, device=current_attention_mask.device)
-                    ], dim=-1)
-            
-            # Final result
-            final_text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
-            
-            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"""