Update app.py

app.py

@@ -1,32 +1,30 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Title
-st.title("Text Generation ")
+# Your model and tokenizer definitions remain the same
+model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
 
-# IMPORTANT: Model Loading Considerations
-# st.warning("Large Language Models (LLMs) like Llama 2-70b can be resource-intensive. This is a simplified demo. Loading the full model might exceed memory limits on standard machines.")
+st.title("Chat with the Language Model")
 
-# Function to load a small portion of the model (for demo purposes)
-@st.cache_resource
-def load_demo_model():
-    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-13b-hf")
-    model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-13b-hf")
-    # ... Add logic to load a tiny slice of the model ... 
-    return tokenizer, model
+# Area to display chat history
+chat_history = st.empty() 
 
-# # Load model on demand 
-# if st.button("Load Demo Model"):
-#     with st.spinner("Loading model (this might still take some time)..."):
-#         
-#         st.success("Demo model loaded!")
+# User input box
+user_input = st.text_input("Your message:", key="input")
 
-# Text input for prompt
-input_text = st.text_area("Enter a prompt for text generation:")
+# Submit button
+if st.button('Send'): 
+    # Add user message to chat history
+    messages.append({"role": "user", "content": user_input})
+    chat_history.text("\n".join([f"**{msg['role']}**: {msg['content']}" for msg in messages]))  
 
-# Button to trigger generation
-if st.button("Generate Text"):
-    # ... Logic to use model.generate() with input_text ...
-    generated_text = "..."  #Replace with your generation code
-    tokenizer, model = load_demo_model()
-    st.write(generated_text)
+    # Process the input with the model
+    inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to("cuda") 
+    outputs = model.generate(inputs, max_new_tokens=20)
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True) 
+
+    # Add model response to chat history
+    messages.append({"role": "assistant", "content": generated_text})
+    chat_history.text("\n".join([f"**{msg['role']}**: {msg['content']}" for msg in messages]))