Update app.py

app.py

@@ -1,34 +1,48 @@
-import spaces 
-import gradio as gr
-from graph_tool import generate_plot
 import os
-import platform
-from dotenv import load_dotenv
-import logging
 import re
 import json
+import time
+import torch
+import gradio as gr
 import threading
+import logging
+import platform
+import warnings
+
 from datetime import datetime
+from dotenv import load_dotenv
 from typing import Annotated, Sequence, TypedDict, List, Optional, Any, Type
+
 from pydantic import BaseModel, Field
 
+# Gradio Spaces decorator (for @spaces.GPU)
+import spaces 
+
 # LangGraph imports
 from langgraph.graph import StateGraph, START, END
 from langgraph.graph.message import add_messages
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.prebuilt import ToolNode
 
-# Updated LangChain imports
+# LangChain Core imports
 from langchain_core.tools import tool
 from langchain_core.messages import HumanMessage, SystemMessage, AIMessage, ToolMessage, BaseMessage
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.runnables import Runnable
 from langchain_core.runnables.utils import Input, Output
 
-from transformers import AutoTokenizer, TextIteratorStreamer, AutoModelForCausalLM, BitsAndBytesConfig
-import torch
-import time
-import warnings
+# Transformers imports
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    TextIteratorStreamer,
+    StoppingCriteria,
+    StoppingCriteriaList,
+    BitsAndBytesConfig,
+)
+
+from graph_tool import generate_plot
+
 
 # Updated environment variables
 os.environ['HF_HOME'] = '/tmp/huggingface'
@@ -405,6 +419,18 @@ class Phi3MiniEducationalLLM(Runnable):
             # Fallback to manual Phi-3 format
             return f"<|system|>\n{SYSTEM_PROMPT}<|end|>\n<|user|>\n{prompt}<|end|>\n<|assistant|>\n"
 
+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
+
+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"""
@@ -432,16 +458,17 @@ class Phi3MiniEducationalLLM(Runnable):
     
             # 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=4096
+                    max_length=max_input_length
                 )
                 
                 # Ensure inputs are properly formatted
-                if not hasattr(inputs, 'input_ids'):
+                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."
                     
@@ -462,7 +489,7 @@ class Phi3MiniEducationalLLM(Runnable):
                     outputs = model.generate(
                         input_ids=inputs['input_ids'],
                         attention_mask=inputs.get('attention_mask', None),
-                        max_new_tokens=1200,
+                        max_new_tokens=300,
                         do_sample=True,
                         temperature=0.7,
                         top_p=0.9,
@@ -470,7 +497,8 @@ class Phi3MiniEducationalLLM(Runnable):
                         repetition_penalty=1.1,
                         pad_token_id=self.tokenizer.eos_token_id,
                         use_cache=False,
-                        past_key_values=None
+                        past_key_values=None,
+                        stopping_criteria=stop_criteria
                     )
                 except Exception as generation_error:
                     logger.error(f"Generation error: {generation_error}")
@@ -480,6 +508,13 @@ class Phi3MiniEducationalLLM(Runnable):
             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."
@@ -489,6 +524,7 @@ class Phi3MiniEducationalLLM(Runnable):
             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 e:
             logger.error(f"Generation error with 4-bit model: {e}")