Update app.py

app.py

@@ -1,34 +1,96 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-# Load your conversational model
-chatbot = pipeline("text-generation", model="eyad-silx/Quasar-32B")  # Your model here
+@st.cache_resource
+def load_model():
+    """Load model and tokenizer with caching"""
+    try:
+        tokenizer = AutoTokenizer.from_pretrained("eyad-silx/Quasar-32B")
+        model = AutoModelForCausalLM.from_pretrained("eyad-silx/Quasar-32B")
+        return model, tokenizer
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        return None, None
 
-# Initialize chat history
-if 'history' not in st.session_state:
-    st.session_state.history = []
-
-# Set up the Streamlit interface
+# Page config
+st.set_page_config(page_title="Chat with Quasar-32B", layout="wide")
 st.title("Chat with Quasar-32B")
 
-# User input
-user_input = st.text_input("You:", "")
+# Initialize session state for chat history
+if 'messages' not in st.session_state:
+    st.session_state.messages = []
 
-# When the user sends a message
-if st.button("Send"):
-    # Add user input to chat history
-    st.session_state.history.append({"role": "user", "content": user_input})
+# Load model and tokenizer
+model, tokenizer = load_model()
 
-    # Generate response from the model
-    response = chatbot(user_input, max_length=50, num_return_sequences=1)[0]['generated_text']
-    
-    # Add model's response to chat history
-    st.session_state.history.append({"role": "assistant", "content": response})
+# Chat interface
+def generate_response(prompt):
+    """Generate response from the model"""
+    try:
+        # Prepare the input
+        inputs = tokenizer(prompt, return_tensors="pt", padding=True)
+        
+        # Generate response
+        with torch.no_grad():
+            outputs = model.generate(
+                inputs["input_ids"],
+                max_length=200,
+                num_return_sequences=1,
+                temperature=0.7,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
+        # Decode and return the response
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response
+    except Exception as e:
+        return f"Error generating response: {str(e)}"
+
+# Chat interface
+st.write("### Chat")
+chat_container = st.container()
 
 # Display chat history
-if st.session_state.history:
-    for chat in st.session_state.history:
-        if chat['role'] == 'user':
-            st.markdown(f"**You:** {chat['content']}")
-        else:
-            st.markdown(f"**AI:** {chat['content']}")
+with chat_container:
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.write(message["content"])
+
+# User input
+if prompt := st.chat_input("Type your message here"):
+    # Add user message to chat history
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    
+    # Display user message
+    with chat_container:
+        with st.chat_message("user"):
+            st.write(prompt)
+    
+    # Generate and display assistant response
+    if model and tokenizer:
+        with st.chat_message("assistant"):
+            with st.spinner("Thinking..."):
+                response = generate_response(prompt)
+                st.write(response)
+                st.session_state.messages.append({"role": "assistant", "content": response})
+    else:
+        st.error("Model failed to load. Please check your configuration.")
+
+# Add a button to clear chat history
+if st.button("Clear Chat History"):
+    st.session_state.messages = []
+    st.experimental_rerun()
+
+# Display system information
+with st.sidebar:
+    st.write("### System Information")
+    st.write("Model: Quasar-32B")
+    st.write("Status: Running" if model and tokenizer else "Status: Not loaded")
+    
+    # Add some helpful instructions
+    st.write("### Instructions")
+    st.write("1. Type your message in the chat input")
+    st.write("2. Press Enter or click Send")
+    st.write("3. Wait for the AI to respond")
+    st.write("4. Use 'Clear Chat History' to start fresh")