Update app.py

app.py

@@ -1,20 +1,28 @@
+# Import necessary libraries
 import streamlit as st
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-# Initialize the pipeline with your chosen model
-pipe = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.2")
+# Load the model and tokenizer
+model_name = "mistralai/Mistral-7B-v0.1"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
 
-# Streamlit app layout
-st.title("Mistral Chatbot")
+# Streamlit app
+def main():
+    st.title("Mistral Chatbot")
+    
+    user_input = st.text_input("You: ", "Hello, chatbot!")
 
-# User input
-user_input = st.text_input("You:", "")
+    if st.button("Send"):
+        with st.spinner("Thinking..."):
+            # Tokenize the user input and generate a response
+            model_inputs = tokenizer(user_input, return_tensors="pt")
+            model_inputs = {k: v.to("cuda") for k, v in model_inputs.items()}
+            generated_ids = model.generate(**model_inputs, max_new_tokens=100, do_sample=True)
+            chatbot_response = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+            
+            st.text_area("Chatbot:", value=chatbot_response, height=200, max_chars=None, key=None)
 
-if st.button("Send"):
-    if user_input:
-        # Generate a response
-        response = pipe(user_input, max_length=50, clean_up_tokenization_spaces=True)
-        # Display the response
-        st.text_area("Chatbot:", value=response[0]['generated_text'], height=100, max_chars=None, key=None)
-    else:
-        st.warning("Please enter some text to chat.")
+if __name__ == "__main__":
+    main()