Update app.py

app.py

@@ -3,28 +3,40 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import os
 from threading import Thread
-import time
 
-# Load Model and Tokenizer
-token = os.environ.get("HF_TOKEN")
-model_name = "large-traversaal/Phi-4-Hindi"
+# Define model path for caching (Avoids reloading every app restart)
+MODEL_PATH = "/mnt/data/Phi-4-Hindi"
+TOKEN = os.environ.get("HF_TOKEN")
+MODEL_NAME = "large-traversaal/Phi-4-Hindi"
 
+# Load Model & Tokenizer Once
 @st.cache_resource()
 def load_model():
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        token=token,
-        trust_remote_code=True,
-        torch_dtype=torch.bfloat16
-    )
-    tok = AutoTokenizer.from_pretrained(model_name, token=token)
-    return model, tok
+    with st.spinner("Loading model... Please wait ⏳"):
+        if not os.path.exists(MODEL_PATH):
+            model = AutoModelForCausalLM.from_pretrained(
+                MODEL_NAME, token=TOKEN, trust_remote_code=True, torch_dtype=torch.bfloat16
+            )
+            tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=TOKEN)
+            model.save_pretrained(MODEL_PATH)
+            tokenizer.save_pretrained(MODEL_PATH)
+        else:
+            model = AutoModelForCausalLM.from_pretrained(MODEL_PATH)
+            tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+    
+    return model, tokenizer
 
+# Load and move model to appropriate device
 model, tok = load_model()
-terminators = [tok.eos_token_id]
-
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = model.to(device)
+try:
+    model = model.to(device)
+except torch.cuda.OutOfMemoryError:
+    st.error("⚠️ CUDA Out of Memory! Running on CPU instead.")
+    device = torch.device("cpu")
+    model = model.to(device)
+
+terminators = [tok.eos_token_id]
 
 # Initialize session state if not set
 if "chat_history" not in st.session_state:
@@ -32,13 +44,19 @@ if "chat_history" not in st.session_state:
 
 # Chat function
 def chat(message, temperature, do_sample, max_tokens):
-    chat_log = st.session_state.chat_history.copy()
-    chat_log.append({"role": "user", "content": message})
-    messages = tok.apply_chat_template(chat_log, tokenize=False, add_generation_prompt=True)
+    """Processes chat input and generates a response using the model."""
     
+    # Append new message to history
+    st.session_state.chat_history.append({"role": "user", "content": message})
+
+    # Convert chat history into model-friendly format
+    messages = tok.apply_chat_template(st.session_state.chat_history, tokenize=False, add_generation_prompt=True)
     model_inputs = tok([messages], return_tensors="pt").to(device)
+    
+    # Initialize streamer for token-wise response
     streamer = TextIteratorStreamer(tok, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
     
+    # Define generation parameters
     generate_kwargs = {
         "inputs": model_inputs["input_ids"],
         "streamer": streamer,
@@ -50,46 +68,61 @@ def chat(message, temperature, do_sample, max_tokens):
     
     if temperature == 0:
         generate_kwargs["do_sample"] = False
-    
+
+    # Generate response asynchronously
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
     
-    partial_text = ""
+    # Collect response as it streams
+    response_text = ""
     for new_text in streamer:
-        partial_text += new_text
-        yield partial_text
-    
-    st.session_state.chat_history.append({"role": "assistant", "content": partial_text})
+        response_text += new_text
+        yield response_text
 
-# Streamlit UI
+    # Save the assistant's response to session history
+    st.session_state.chat_history.append({"role": "assistant", "content": response_text})
+
+# UI Setup
 st.title("💬 Chat With Phi-4-Hindi")
+st.success("✅ Model is READY to chat!")
 st.markdown("Chat with [large-traversaal/Phi-4-Hindi](https://huggingface.co/large-traversaal/Phi-4-Hindi)")
 
-# Chat input
+# Sidebar Chat Settings
 temperature = st.sidebar.slider("Temperature", 0.0, 1.0, 0.3, 0.1)
 do_sample = st.sidebar.checkbox("Use Sampling", value=True)
 max_tokens = st.sidebar.slider("Max Tokens", 128, 4096, 512, 1)
 text_color = st.sidebar.selectbox("Text Color", ["Red", "Black", "Blue", "Green", "Purple"], index=0)
 dark_mode = st.sidebar.checkbox("🌙 Dark Mode", value=False)
 
+# Function to format chat messages
 def get_html_text(text, color):
     return f'<p style="color: {color.lower()}; font-size: 16px;">{text}</p>'
 
+# Display chat history
 for msg in st.session_state.chat_history:
-    if msg["role"] == "user":
-        st.markdown(get_html_text("👤 " + msg["content"], "black"), unsafe_allow_html=True)
-    else:
-        st.markdown(get_html_text("🤖 " + msg["content"], text_color), unsafe_allow_html=True)
+    role = "👤" if msg["role"] == "user" else "🤖"
+    st.markdown(get_html_text(f"**{role}:** {msg['content']}", text_color if role == "🤖" else "black"), unsafe_allow_html=True)
 
+# User Input Handling
 user_input = st.text_input("Type your message:", "")
+
 if st.button("Send"):
     if user_input.strip():
         st.session_state.chat_history.append({"role": "user", "content": user_input})
-        with st.spinner("Generating response..."):
-            for output in chat(user_input, temperature, do_sample, max_tokens):
-                pass
+        
+        # Display chatbot response
+        with st.spinner("Generating response... 🤖💭"):
+            response_generator = chat(user_input, temperature, do_sample, max_tokens)
+            final_response = ""
+            for output in response_generator:
+                final_response = output  # Store latest output
+        
+        #st.success("✅ Response generated!")
+        # Add generated response to session state
         st.experimental_rerun()
 
 if st.button("🧹 Clear Chat"):
-    st.session_state.chat_history = []
+    with st.spinner("Clearing chat history..."):
+        st.session_state.chat_history = []
+    st.success("✅ Chat history cleared!")
     st.experimental_rerun()