Update app.py

app.py

@@ -1,25 +1,35 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Load the model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b")
-model = AutoModelForCausalLM.from_pretrained("google/gemma-7b", device="cuda" if st.sidebar.checkbox("Use GPU", True) else "cpu")
+model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, device="cuda" if st.sidebar.checkbox("Use GPU", True) else "cpu")
 
-# Function to generate text based on user input
-def generate_text(prompt):
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
-    outputs = model.generate(input_ids, max_length=100)
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return generated_text
+# Function to generate responses based on user messages
+def generate_response(messages):
+    inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)
+    outputs = model.generate(inputs, max_new_tokens=100)
+    generated_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return generated_response
 
 # Streamlit app
-st.title("Text Generation with Google Gemma 7b")
+st.title("Mixtral Chatbot")
 
-prompt = st.text_area("Enter your prompt here:", "")
-if st.button("Generate Text"):
-    if prompt:
-        generated_text = generate_text(prompt)
-        st.write("Generated Text:")
-        st.write(generated_text)
+messages = []
+user_input = st.text_input("You:", "")
+
+if st.button("Send"):
+    if user_input:
+        messages.append({"role": "user", "content": user_input})
+        bot_response = generate_response(messages)
+        messages.append({"role": "assistant", "content": bot_response})
     else:
-        st.warning("Please enter a prompt.")
+        st.warning("Please enter a message.")
+
+# Display conversation
+for message in messages:
+    if message["role"] == "user":
+        st.text_input("You:", value=message["content"], disabled=True)
+    elif message["role"] == "assistant":
+        st.text_area("Mixtral:", value=message["content"], disabled=True)