Update app.py

app.py

@@ -415,85 +415,106 @@ class Phi3MiniEducationalLLM(Runnable):
             return f"[Error generating response: {str(e)}]"
 
     def stream_generate(self, input: Input, config=None):
-        """Streaming generation using TextIteratorStreamer"""
-        start_stream_time = time.perf_counter()
-        current_time = datetime.now()
-        logger.info("Starting stream_generate with TextIteratorStreamer...")
-        
-        # Handle both string and dict inputs
-        if isinstance(input, dict):
-            prompt = input.get('input', str(input))
-        else:
-            prompt = str(input)
-            
-        try:
-            # Format using Phi-3 chat template
-            text = self._format_chat_template(prompt)
+    """Streaming generation using TextIteratorStreamer with loop detection and early escape."""
+    start_stream_time = time.perf_counter()
+    current_time = datetime.now()
+    logger.info("Starting stream_generate with TextIteratorStreamer and loop detection...")
 
-            inputs = self.tokenizer(
-                text, 
-                return_tensors="pt", 
-                padding=True, 
-                truncation=True, 
-                max_length=3072
-            )
+    if isinstance(input, dict):
+        prompt = input.get('input', str(input))
+    else:
+        prompt = str(input)
 
-            # Move inputs to model device
-            inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
+    try:
+        # Format using Phi-3 chat template
+        text = self._format_chat_template(prompt)
+
+        inputs = self.tokenizer(
+            text,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=3072
+        )
 
-            # Initialize TextIteratorStreamer
-            streamer = TextIteratorStreamer(
-                self.tokenizer, 
-                skip_prompt=True, 
-                skip_special_tokens=True
-            )
+        # Move inputs to model device
+        inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
 
-            # Generation parameters
-            generation_kwargs = {
-                **inputs,
-                "max_new_tokens": 800,
-                "do_sample": True,
-                "temperature": 0.7,
-                "top_p": 0.9,
-                "top_k": 50,
-                "repetition_penalty": 1.1,
-                "pad_token_id": self.tokenizer.eos_token_id,
-                "streamer": streamer,
-                "use_cache": True
-            }
+        # Initialize TextIteratorStreamer
+        streamer = TextIteratorStreamer(
+            self.tokenizer,
+            skip_prompt=True,
+            skip_special_tokens=True
+        )
 
-            # Start generation in a separate thread
-            generation_thread = threading.Thread(
-                target=self.model.generate, 
-                kwargs=generation_kwargs
-            )
-            generation_thread.start()
+        # Generation parameters
+        generation_kwargs = {
+            **inputs,
+            "max_new_tokens": 800,
+            "do_sample": True,
+            "temperature": 0.7,
+            "top_p": 0.9,
+            "top_k": 50,
+            "repetition_penalty": 1.2,  # Slightly stronger to help with loop prevention
+            "pad_token_id": self.tokenizer.eos_token_id,
+            "streamer": streamer,
+            "use_cache": True
+        }
+
+        # Start generation in background
+        generation_thread = threading.Thread(
+            target=self.model.generate,
+            kwargs=generation_kwargs
+        )
+        generation_thread.start()
 
-            # Yield tokens as they become available
-            generated_text = ""
-            try:
-                for new_text in streamer:
-                    if new_text:  # Only yield non-empty strings
-                        generated_text += new_text
-                        yield generated_text
-            except Exception as e:
-                logger.error(f"Error in streaming iteration: {e}")
-                yield f"[Streaming error: {str(e)}]"
+        # Track outputs
+        generated_text = ""
+        token_history = []
+        loop_window = 20  # Number of tokens to compare
+        loop_threshold = 3  # Allow N repetitions before aborting
 
-            # Wait for generation to complete
-            generation_thread.join()
-            
-            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. Generated length: {len(generated_text)} chars. Model: {self.model_name}. Timestamp: {current_time:%Y-%m-%d %H:%M:%S}")
-            logger.info(f"Stream generation completed: {len(generated_text)} chars in {stream_time:.2f}s")
-            
+        try:
+            for new_text in streamer:
+                if not new_text:
+                    continue
+
+                generated_text += new_text
+
+                # Tokenize and track
+                tokens = self.tokenizer.tokenize(new_text)
+                token_history.extend(tokens)
+
+                # Check for loops
+                if len(token_history) >= 2 * loop_window:
+                    recent = token_history[-loop_window:]
+                    prev = token_history[-2*loop_window:-loop_window]
+                    overlap = sum(1 for r, p in zip(recent, prev) if r == p)
+
+                    if overlap >= loop_threshold:
+                        logger.warning(f"Looping detected (overlap: {overlap}/{loop_window}). Aborting generation.")
+                        yield "[Looping detected — generation stopped early]"
+                        break
+
+                yield generated_text
         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)}]"
+            logger.error(f"Error in streaming iteration: {e}")
+            yield f"[Streaming error: {str(e)}]"
+
+        generation_thread.join()
+
+        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. Generated length: {len(generated_text)} chars. Model: {self.model_name}. Timestamp: {current_time:%Y-%m-%d %H:%M:%S}")
+        logger.info(f"Stream generation completed: {len(generated_text)} chars in {stream_time:.2f}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)}]"
+
 
     @property
     def InputType(self) -> Type[Input]: