Update app.py

app.py

@@ -336,7 +336,7 @@ Rather than providing complete solutions, you should:
 
 Your goal is to be an educational partner who empowers students to succeed through understanding."""
 
-# FIXED LLM Class with Phi-3-mini
+# --- LLM Class with Phi-3 Mini ---
 class Phi3MiniEducationalLLM(Runnable):
     """LLM class optimized for Microsoft Phi-3-mini-4k-instruct with 4-bit quantization"""
     
@@ -419,243 +419,231 @@ 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
-        self.stop_sequence = tokenizer.encode(stop_sequence, add_special_tokens=False)
-
-    def __call__(self, input_ids, scores, **kwargs):
-        if input_ids[0, -len(self.stop_sequence):].tolist() == self.stop_sequence:
-            return True
-        return False
-
-@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()
 
-    # Handle different input types
-    if isinstance(input, dict):
-        if 'input' in input:
-            prompt = input['input']
-        elif 'messages' in input:
-            prompt = str(input['messages'])
+        # 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)
-    else:
-        prompt = str(input)
-
-    try:
-        model = self._load_model_if_needed()
-        text = self._format_chat_template(prompt)
 
         try:
-            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")
+            model = self._load_model_if_needed()
+            text = self._format_chat_template(prompt)
+
+            try:
+                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."
+            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 tokenizer_error:
-            logger.error(f"Tokenization error: {tokenizer_error}")
-            return "I encountered an error processing your request. Please try again."
 
-        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}")
-            return "I encountered an error processing your request. Please try again."
+            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}")
+                return "I encountered an error processing your request. Please try again."
 
-        # Define stopping criteria after tokenizer initialization
-        stop_criteria = StoppingCriteriaList([StopOnSequence(self.tokenizer, "User:")])
+            # Define stopping criteria after tokenizer initialization
+            stop_criteria = StoppingCriteriaList([StopOnSequence(self.tokenizer, "User:")])
+
+            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."
 
-        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."
+                new_tokens = outputs[0][len(inputs['input_ids'][0]):]
+                result = self.tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
 
-        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}"
+            )
 
-            # 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."
 
-        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)}"
 
-    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...")
 
-@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)
 
-    # 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:
-        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
+            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 = self.tokenizer(
-                text,
-                return_tensors="pt",
-                padding=True,
-                truncation=True,
-                max_length=4096
-            )
-            if 'input_ids' not in inputs:
+            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
-        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
-        )
+            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()
+            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
+            }
 
-        generated_text = ""
-        consecutive_repeats = 0
-        last_chunk = ""
+            generation_thread = threading.Thread(
+                target=model.generate,
+                kwargs=generation_kwargs
+            )
+            generation_thread.start()
 
-        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
+            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
+            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
-        except Exception as e:
-            logger.error(f"Error in streaming iteration: {e}")
+
+            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
-
-        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."
+                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 InputType(self) -> Type[Input]:
+        return str
 
-@property 
-def OutputType(self) -> Type[Output]:
-    return str
+    @property 
+    def OutputType(self) -> Type[Output]:
+        return str
 
 # LangGraph Agent Implementation with Tool Calling
 class Educational_Agent: