Update app.py

app.py

@@ -1,23 +1,30 @@
+import os
 import streamlit as st
 from transformers import pipeline
 
-# Initialize the pipeline with the Mistral model
-generator = pipeline(model="bigscience/Mistral", model_kwargs={"use_auth_token": True})
+# Get the Hugging Face API token from the environment variable
+HF_API_TOKEN = os.getenv("HF_API_TOKEN")
 
-# Streamlit app setup
-st.title('Mistral LLM Explorer with Transformers')
-st.write('Enter your prompt below to interact with the Mistral model.')
+if not HF_API_TOKEN:
+    st.error("Hugging Face API token not found. Please set the HF_API_TOKEN environment variable.")
+else:
+    # Initialize the pipeline with the Mistral model and your API token
+    generator = pipeline(model="bigscience/Mistral", model_kwargs={"use_auth_token": HF_API_TOKEN})
 
-# User input
-user_input = st.text_area("Prompt:")
+    # Streamlit app setup
+    st.title('Mistral LLM Explorer with Transformers')
+    st.write('Enter your prompt below to interact with the Mistral model.')
 
-# When the user presses the 'Submit' button
-if st.button('Submit'):
-    if user_input:
-        with st.spinner('Generating response...'):
-            # Generate response from the Mistral model
-            results = generator(user_input, max_length=100, clean_up_tokenization_spaces=True)
-            # Display the generated text
-            st.text_area("Response:", value=results[0]['generated_text'], height=300)
-    else:
-        st.warning('Please enter a prompt.')
+    # User input
+    user_input = st.text_area("Prompt:")
+
+    # When the user presses the 'Submit' button
+    if st.button('Submit'):
+        if user_input:
+            with st.spinner('Generating response...'):
+                # Generate response from the Mistral model
+                results = generator(user_input, max_length=100, clean_up_tokenization_spaces=True)
+                # Display the generated text
+                st.text_area("Response:", value=results[0]['generated_text'], height=300)
+        else:
+            st.warning('Please enter a prompt.')