Update app.py

app.py

@@ -419,6 +419,8 @@ class Phi3MiniEducationalLLM(Runnable):
             # Fallback to manual Phi-3 format
             return f"<|system|>\n{SYSTEM_PROMPT}<|end|>\n<|user|>\n{prompt}<|end|>\n<|assistant|>\n"
 
+from transformers import StoppingCriteria, StoppingCriteriaList
+
 class StopOnSequence(StoppingCriteria):
     def __init__(self, tokenizer, stop_sequence):
         self.tokenizer = tokenizer
@@ -429,251 +431,231 @@ class StopOnSequence(StoppingCriteria):
             return True
         return False
 
-stop_criteria = StoppingCriteriaList([StopOnSequence(self.tokenizer, "User:")])
+@spaces.GPU(duration=180)
+def invoke(self, input: Input, config=None) -> Output:
+    """Main invoke method optimized for 4-bit quantized Phi‑3‑mini"""
+    start_invoke_time = time.perf_counter()
+    current_time = datetime.now()
 
-    @spaces.GPU(duration=180)
-    def invoke(self, input: Input, config=None) -> Output:
-        """Main invoke method optimized for 4-bit quantized Phi-3-mini"""
-        start_invoke_time = time.perf_counter()
-        current_time = datetime.now()
-        
-        # FIX: Handle different input types properly
-        if isinstance(input, dict):
-            if 'input' in input:
-                prompt = input['input']
-            elif 'messages' in input:
-                # Handle messages format
-                prompt = str(input['messages'])
-            else:
-                prompt = str(input)
+    # Handle different input types
+    if isinstance(input, dict):
+        if 'input' in input:
+            prompt = input['input']
+        elif 'messages' in input:
+            prompt = str(input['messages'])
         else:
             prompt = str(input)
-            
+    else:
+        prompt = str(input)
+
+    try:
+        model = self._load_model_if_needed()
+        text = self._format_chat_template(prompt)
+
         try:
-            # Load model inside GPU context
-            model = self._load_model_if_needed()
-            
-            # Format using Phi-3 chat template
-            text = self._format_chat_template(prompt)
-    
-            # FIX: Proper tokenization with error handling
-            try:
-                max_input_length = 4096 - 300
-                inputs = self.tokenizer(
-                    text, 
-                    return_tensors="pt", 
-                    padding=True, 
-                    truncation=True, 
-                    max_length=max_input_length
-                )
-                
-                # Ensure inputs are properly formatted
-                if 'input_ids' not in inputs:
-                    logger.error("Tokenizer did not return input_ids")
-                    return "I encountered an error processing your request. Please try again."
-                    
-            except Exception as tokenizer_error:
-                logger.error(f"Tokenization error: {tokenizer_error}")
-                return "I encountered an error processing your request. Please try again."
-    
-            # Move inputs to model device
-            try:
-                inputs = {k: v.to(model.device) for k, v in inputs.items()}
-            except Exception as device_error:
-                logger.error(f"Device transfer error: {device_error}")
+            max_input_length = 4096 - 300
+            inputs = self.tokenizer(
+                text,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=max_input_length
+            )
+            if 'input_ids' not in inputs:
+                logger.error("Tokenizer did not return input_ids")
                 return "I encountered an error processing your request. Please try again."
-    
-            # Generate with optimized parameters for quantized model
-            with torch.no_grad():
-                try:
-                    outputs = model.generate(
-                        input_ids=inputs['input_ids'],
-                        attention_mask=inputs.get('attention_mask', None),
-                        max_new_tokens=300,
-                        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,
-                        use_cache=False,
-                        past_key_values=None,
-                        stopping_criteria=stop_criteria
-                    )
-                except Exception as generation_error:
-                    logger.error(f"Generation error: {generation_error}")
-                    return "I encountered an error generating the response. Please try again."
-    
-            # Decode only new tokens
-            try:
-                new_tokens = outputs[0][len(inputs['input_ids'][0]):]
-                result = self.tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
-            
-                # Soft stop cleanup
-                for stop_word in ["User:", "\n\n", "###"]:
-                    if stop_word in result:
-                        result = result.split(stop_word)[0].strip()
-                        break
-            
-            except Exception as decode_error:
-                logger.error(f"Decoding error: {decode_error}")
-                return "I encountered an error processing the response. Please try again."
-            
-            end_invoke_time = time.perf_counter()
-            invoke_time = end_invoke_time - start_invoke_time
-            log_metric(f"LLM Invoke time (4-bit): {invoke_time:0.4f} seconds. Input length: {len(prompt)} chars. Model: {self.model_name}. Timestamp: {current_time:%Y-%m-%d %H:%M:%S}")
-            
-            return result if result else "I'm still learning how to respond to that properly."
+        except Exception as tokenizer_error:
+            logger.error(f"Tokenization error: {tokenizer_error}")
+            return "I encountered an error processing your request. Please try again."
 
-    
-        except Exception as e:
-            logger.error(f"Generation error with 4-bit model: {e}")
-            end_invoke_time = time.perf_counter()
-            invoke_time = end_invoke_time - start_invoke_time
-            log_metric(f"LLM Invoke time (error): {invoke_time:0.4f} seconds. Model: {self.model_name}. Timestamp: {current_time:%Y-%m-%d %H:%M:%S}")
-            return f"I encountered an error: {str(e)}"
-
-    @spaces.GPU(duration=240)
-    def stream_generate(self, input: Input, config=None):
-        """Streaming generation with 4-bit quantized model and expanded context"""
-        start_stream_time = time.perf_counter()
-        current_time = datetime.now()
-        logger.info("Starting stream_generate with 4-bit quantized model...")
-    
-        # Handle input properly
-        if isinstance(input, dict):
-            if 'input' in input:
-                prompt = input['input']
-            else:
-                prompt = str(input)
-        else:
-            prompt = str(input)
-    
         try:
-            # Load quantized model inside GPU context
-            model = self._load_model_if_needed()
-        
-            # Clear GPU cache
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-            
-            text = self._format_chat_template(prompt)
-    
-            # Proper tokenization with error handling
+            inputs = {k: v.to(model.device) for k, v in inputs.items()}
+        except Exception as device_error:
+            logger.error(f"Device transfer error: {device_error}")
+            return "I encountered an error processing your request. Please try again."
+
+        # Define stopping criteria after tokenizer initialization
+        stop_criteria = StoppingCriteriaList([StopOnSequence(self.tokenizer, "User:")])
+
+        with torch.no_grad():
             try:
-                inputs = self.tokenizer(
-                    text,
-                    return_tensors="pt",
-                    padding=True,
-                    truncation=True,
-                    max_length=4096
+                outputs = model.generate(
+                    input_ids=inputs['input_ids'],
+                    attention_mask=inputs.get('attention_mask', None),
+                    max_new_tokens=300,
+                    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,
+                    use_cache=False,
+                    past_key_values=None,
+                    stopping_criteria=stop_criteria
                 )
-                
-                if not hasattr(inputs, 'input_ids'):
-                    yield "I encountered an error processing your request. Please try again."
-                    return
-                    
-            except Exception as tokenizer_error:
-                logger.error(f"Streaming tokenization error: {tokenizer_error}")
-                yield "I encountered an error processing your request. Please try again."
-                return
-    
-            # Move inputs to model device
-            try:
-                inputs = {k: v.to(model.device) for k, v in inputs.items()}
-            except Exception as device_error:
-                logger.error(f"Streaming device transfer error: {device_error}")
+            except Exception as generation_error:
+                logger.error(f"Generation error: {generation_error}")
+                return "I encountered an error generating the response. Please try again."
+
+        try:
+            new_tokens = outputs[0][len(inputs['input_ids'][0]):]
+            result = self.tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
+
+            # Apply soft-stop cleanup
+            for stop_word in ["User:", "\n\n", "###"]:
+                if stop_word in result:
+                    result = result.split(stop_word)[0].strip()
+                    break
+        except Exception as decode_error:
+            logger.error(f"Decoding error: {decode_error}")
+            return "I encountered an error processing the response. Please try again."
+
+        end_invoke_time = time.perf_counter()
+        invoke_time = end_invoke_time - start_invoke_time
+        log_metric(
+            f"LLM Invoke time (4‑bit): {invoke_time:0.4f} seconds. "
+            f"Input length: {len(prompt)} chars. "
+            f"Model: {self.model_name}. "
+            f"Timestamp: {current_time:%Y‑%m‑%d %H:%M:%S}"
+        )
+
+        return result if result else "I'm still learning how to respond to that properly."
+
+    except Exception as e:
+        logger.error(f"Generation error with 4‑bit model: {e}")
+        end_invoke_time = time.perf_counter()
+        invoke_time = end_invoke_time - start_invoke_time
+        log_metric(
+            f"LLM Invoke time (error): {invoke_time:0.4f} seconds. "
+            f"Model: {self.model_name}. "
+            f"Timestamp: {current_time:%Y‑%m‑%d %H:%M:%S}"
+        )
+        return f"I encountered an error: {str(e)}"
+
+@spaces.GPU(duration=240)
+def stream_generate(self, input: Input, config=None):
+    """Streaming generation with 4‑bit quantized model and expanded context"""
+    start_stream_time = time.perf_counter()
+    current_time = datetime.now()
+    logger.info("Starting stream_generate with 4‑bit quantized model...")
+
+    # Handle input properly
+    if isinstance(input, dict):
+        prompt = input.get('input', str(input))
+    else:
+        prompt = str(input)
+
+    try:
+        model = self._load_model_if_needed()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        text = self._format_chat_template(prompt)
+
+        try:
+            inputs = self.tokenizer(
+                text,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=4096
+            )
+            if 'input_ids' not in inputs:
                 yield "I encountered an error processing your request. Please try again."
                 return
-    
-            # Initialize TextIteratorStreamer - this streams the GENERATED TOKENS, not the input
-            streamer = TextIteratorStreamer(
-                self.tokenizer,
-                skip_prompt=True,  # Skip the input prompt in output
-                skip_special_tokens=True
-            )
-    
-            # Generation parameters optimized for 4-bit
-            generation_kwargs = {
-                "input_ids": inputs['input_ids'],
-                "attention_mask": inputs.get('attention_mask', None),
-                "max_new_tokens": 1200,
-                "do_sample": True,
-                "temperature": 0.7,
-                "top_p": 0.9,
-                "top_k": 50,
-                "repetition_penalty": 1.2,
-                "pad_token_id": self.tokenizer.eos_token_id,
-                "streamer": streamer,  # This streams the OUTPUT tokens as they're generated
-                "use_cache": False,
-                "past_key_values": None
-            }
-    
-            # Start generation in background thread
-            generation_thread = threading.Thread(
-                target=model.generate,
-                kwargs=generation_kwargs
-            )
-            generation_thread.start()
-    
-            # Stream the generated tokens as they come from the model
-            generated_text = ""
-            consecutive_repeats = 0
-            last_chunk = ""
-    
-            try:
-                # This loop receives tokens as they're generated by the model
-                for new_token_text in streamer:
-                    if not new_token_text:
-                        continue
-    
-                    # Accumulate the generated text
-                    generated_text += new_token_text
-    
-                    # Simple repetition detection
-                    if new_token_text == last_chunk:
-                        consecutive_repeats += 1
-                        if consecutive_repeats >= 5:
-                            logger.warning("Repetitive generation detected, stopping early")
-                            break
-                    else:
-                        consecutive_repeats = 0
-                        last_chunk = new_token_text
-    
-                    # Yield the accumulated generated text (not the input prompt)
-                    yield generated_text
-                    
-            except Exception as e:
-                logger.error(f"Error in streaming iteration: {e}")
-                if not generated_text.strip():
-                    generated_text = "I apologize, but I'm having trouble generating a response. Please try rephrasing your question."
+        except Exception as tokenizer_error:
+            logger.error(f"Streaming tokenization error: {tokenizer_error}")
+            yield "I encountered an error processing your request. Please try again."
+            return
+
+        try:
+            inputs = {k: v.to(model.device) for k, v in inputs.items()}
+        except Exception as device_error:
+            logger.error(f"Streaming device transfer error: {device_error}")
+            yield "I encountered an error processing your request. Please try again."
+            return
+
+        streamer = TextIteratorStreamer(
+            self.tokenizer,
+            skip_prompt=True,
+            skip_special_tokens=True
+        )
+
+        generation_kwargs = {
+            "input_ids": inputs['input_ids'],
+            "attention_mask": inputs.get('attention_mask', None),
+            "max_new_tokens": 1200,
+            "do_sample": True,
+            "temperature": 0.7,
+            "top_p": 0.9,
+            "top_k": 50,
+            "repetition_penalty": 1.2,
+            "pad_token_id": self.tokenizer.eos_token_id,
+            "streamer": streamer,
+            "use_cache": False,
+            "past_key_values": None
+        }
+
+        generation_thread = threading.Thread(
+            target=model.generate,
+            kwargs=generation_kwargs
+        )
+        generation_thread.start()
+
+        generated_text = ""
+        consecutive_repeats = 0
+        last_chunk = ""
+
+        try:
+            for new_token_text in streamer:
+                if not new_token_text:
+                    continue
+                generated_text += new_token_text
+                if new_token_text == last_chunk:
+                    consecutive_repeats += 1
+                    if consecutive_repeats >= 5:
+                        logger.warning("Repetitive generation detected, stopping early")
+                        break
+                else:
+                    consecutive_repeats = 0
+                    last_chunk = new_token_text
                 yield generated_text
-    
-            generation_thread.join()
-    
-            # Ensure we have some content
+        except Exception as e:
+            logger.error(f"Error in streaming iteration: {e}")
             if not generated_text.strip():
                 generated_text = "I apologize, but I'm having trouble generating a response. Please try rephrasing your question."
-                yield generated_text
-    
-            end_stream_time = time.perf_counter()
-            stream_time = end_stream_time - start_stream_time
-            log_metric(f"LLM Stream time (4-bit): {stream_time:0.4f} seconds. Generated length: {len(generated_text)} chars. Model: {self.model_name}. Timestamp: {current_time:%Y-%m-%d %H:%M:%S}")
-    
-        except Exception as e:
-            logger.error(f"4-bit 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 "I encountered an error generating the response. Please try again."
-
-    @property
-    def InputType(self) -> Type[Input]:
-        return str
-
-    @property 
-    def OutputType(self) -> Type[Output]:
-        return str
+            yield generated_text
+
+        generation_thread.join()
+        if not generated_text.strip():
+            generated_text = "I apologize, but I'm having trouble generating a response. Please try rephrasing your question."
+            yield generated_text
+
+        end_stream_time = time.perf_counter()
+        stream_time = end_stream_time - start_stream_time
+        log_metric(
+            f"LLM Stream time (4‑bit): {stream_time:0.4f} seconds. "
+            f"Generated length: {len(generated_text)} chars. "
+            f"Model: {self.model_name}. "
+            f"Timestamp: {current_time:%Y‑%m‑%d %H:%M:%S}"
+        )
+    except Exception as e:
+        logger.error(f"4‑bit 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. "
+            f"Model: {self.model_name}. "
+            f"Timestamp: {current_time:%Y‑%m‑%d %H:%M:%S}"
+        )
+        yield "I encountered an error generating the response. Please try again."
+
+@property
+def InputType(self) -> Type[Input]:
+    return str
+
+@property 
+def OutputType(self) -> Type[Output]:
+    return str
 
 # LangGraph Agent Implementation with Tool Calling
 class Educational_Agent: