Update app.py

app.py

@@ -1,93 +1,137 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
 
-# Custom CSS for styling
+# ======================
+# 🎨 STYLING & LAYOUT
+# ======================
+st.set_page_config(
+    page_title="LLM Fine-Tuning Studio",
+    page_icon="🧠",
+    layout="wide"
+)
+
+# Custom CSS
 st.markdown("""
-    <style>
+<style>
+    /* Main container */
     .main {
-        background-color: #f5f5f5;
+        background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
     }
-    .stTextInput>div>div>input {
-        background-color: #ffffff;
-        color: #000000;
+    
+    /* Headers */
+    h1 {
+        color: #2c3e50;
+        text-align: center;
+        font-family: 'Arial', sans-serif;
+        border-bottom: 2px solid #4CAF50;
+        padding-bottom: 10px;
+    }
+    
+    /* Sidebar */
+    [data-testid="stSidebar"] {
+        background: linear-gradient(195deg, #2c3e50 0%, #4CAF50 100%) !important;
+        color: white;
     }
+    
+    /* Buttons */
     .stButton>button {
-        background-color: #4CAF50;
+        background: linear-gradient(to right, #4CAF50, #2E8B57);
         color: white;
-        border-radius: 5px;
         border: none;
+        border-radius: 25px;
         padding: 10px 24px;
+        font-weight: bold;
     }
-    .stButton>button:hover {
-        background-color: #45a049;
+    
+    /* Chat bubbles */
+    .user-message {
+        background: #e3f2fd;
+        border-radius: 15px 15px 0 15px;
+        padding: 12px;
+        margin: 5px 0;
     }
-    .title {
-        font-size: 2.5em;
-        color: #2c3e50;
-        text-align: center;
-        margin-bottom: 0.5em;
-    }
-    .sidebar .sidebar-content {
-        background-color: #2c3e50;
+    
+    .bot-message {
+        background: #4CAF50;
         color: white;
+        border-radius: 15px 15px 15px 0;
+        padding: 12px;
+        margin: 5px 0;
     }
-    </style>
-    """, unsafe_allow_html=True)
-
-# App Title
-st.markdown('<p class="title">💬 Fine-Tuned LLM Chat</p>', unsafe_allow_html=True)
-
-# Sidebar for settings
-with st.sidebar:
-    st.header("⚙️ Settings")
-    model_name = st.selectbox(
-        "Select Model",
-        ["mistralai/Mistral-7B-v0.1", "meta-llama/Llama-2-7b-chat-hf"],
-        help="Choose a pre-trained model to fine-tune."
-    )
-    temperature = st.slider(
-        "Temperature", 
-        min_value=0.1, 
-        max_value=1.0, 
-        value=0.7,
-        help="Controls randomness (lower = more deterministic)."
-    )
-    max_length = st.slider(
-        "Max Response Length", 
-        min_value=50, 
-        max_value=500, 
-        value=150,
-        help="Maximum number of tokens in the response."
-    )
+    
+    /* Input box */
+    .stTextInput>div>div>input {
+        border-radius: 20px !important;
+        padding: 10px 15px !important;
+    }
+</style>
+""", unsafe_allow_html=True)
 
-# Load model (cached to avoid reloading)
+# ======================
+# 🧠 MODEL LOADING
+# ======================
 @st.cache_resource
-def load_model(model_name):
+def load_model():
+    model_name = "mistralai/Mistral-7B-v0.1"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float16,
+        device_map="auto"
+    )
     return tokenizer, model
 
-tokenizer, model = load_model(model_name)
+tokenizer, model = load_model()
+
+# ======================
+# 🎛️ SIDEBAR CONTROLS
+# ======================
+with st.sidebar:
+    st.title("⚙️ Fine-Tuning Controls")
+    
+    st.subheader("Model Parameters")
+    temperature = st.slider("Temperature", 0.1, 1.0, 0.7, 0.05)
+    max_length = st.slider("Max Length", 50, 500, 150)
+    
+    st.subheader("Fine-Tuning Options")
+    use_lora = st.checkbox("Use LoRA", True)
+    quantize = st.selectbox("Quantization", ["None", "4-bit", "8-bit"])
+    
+    if st.button("🔄 Apply Changes"):
+        st.toast("Settings updated!", icon="✅")
+
+# ======================
+# 💬 MAIN CHAT INTERFACE
+# ======================
+st.title("🧠 LLM Fine-Tuning Studio")
+st.caption("Fine-tune and deploy state-of-the-art language models")
 
 # Initialize chat history
 if "messages" not in st.session_state:
-    st.session_state.messages = []
+    st.session_state.messages = [
+        {"role": "assistant", "content": "Hello! I'm your fine-tuned AI assistant. How can I help you today?"}
+    ]
 
 # Display chat messages
 for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
+    if message["role"] == "assistant":
+        with st.chat_message("assistant"):
+            st.markdown(f'<div class="bot-message">{message["content"]}</div>', unsafe_allow_html=True)
+    else:
+        with st.chat_message("user"):
+            st.markdown(f'<div class="user-message">{message["content"]}</div>', unsafe_allow_html=True)
 
 # Chat input
-if prompt := st.chat_input("Ask me anything..."):
+if prompt := st.chat_input("Type your message..."):
     st.session_state.messages.append({"role": "user", "content": prompt})
+    
     with st.chat_message("user"):
-        st.markdown(prompt)
-
+        st.markdown(f'<div class="user-message">{prompt}</div>', unsafe_allow_html=True)
+    
     # Generate response
     with st.chat_message("assistant"):
-        with st.spinner("Thinking..."):
+        with st.spinner("🧠 Thinking..."):
             inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
             outputs = model.generate(
                 **inputs,
@@ -96,5 +140,15 @@ if prompt := st.chat_input("Ask me anything..."):
                 do_sample=True
             )
             response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-            st.markdown(response)
-    st.session_state.messages.append({"role": "assistant", "content": response})
+            
+            st.markdown(f'<div class="bot-message">{response}</div>', unsafe_allow_html=True)
+            st.session_state.messages.append({"role": "assistant", "content": response})
+
+# ======================
+# 📊 FINE-TUNE STATUS
+# ======================
+st.sidebar.markdown("---")
+st.sidebar.subheader("Training Metrics")
+st.sidebar.metric("Loss", "0.45", delta="-0.02")
+st.sidebar.metric("Accuracy", "87%", delta="+2%")
+st.sidebar.progress(75)