Update app.py

app.py

@@ -1,26 +1,29 @@
+
+import streamlit as st
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import os
-from transformers import pipeline
-# Retrieve the API token from environment variables
-api_token = os.getenv('hugging_face_api_token')
 
+# Check and retrieve the API token from environment variables
+api_token = os.getenv('hugging_face_api_token')
 if not api_token:
-    raise ValueError("Hugging Face API token not found. Please set the HUGGING_FACE_API_TOKEN environment variable.")
+    st.error("Hugging Face API token not found. Please set the HUGGING_FACE_API_TOKEN environment variable.")
+    st.stop()
 
 # Configure the use of the token for Hugging Face operations
 from huggingface_hub import HfFolder
 HfFolder.save_token(api_token)
 
-
-generator = pipeline("text-generation", model="google/gemma-7b", tokenizer="google/gemma-7b")
-
-
+# Initialize tokenizer and model with the correct model ID
+model_id = "mistral-community/Mistral-8x22B-v0.1"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id)
 
 # Streamlit app title and description
-st.title("Gemma Text Generation App")
-st.write("This app generates text based on the input prompt using the Gemma model.")
+st.title("Text Generation App")
+st.write("This app generates text based on the input prompt using the Mistral-8x22B model.")
 
 # Text input for user prompt
-prompt = st.text_input("Enter your prompt:", "Once upon a time,")
+prompt = st.text_input("Enter your prompt:", "Hello my name is")
 
 # User controls for output length and creativity
 max_length = st.slider("Select the maximum output length:", min_value=50, max_value=500, value=100)
@@ -29,8 +32,10 @@ temperature = st.slider("Adjust the creativity level (temperature):", min_value=
 # Generate button to trigger text generation
 if st.button("Generate Text"):
     with st.spinner('Generating text...'):
+        inputs = tokenizer(prompt, return_tensors="pt")
         try:
-            generated_text = generator(prompt, max_length=max_length, temperature=temperature)[0]['generated_text']
+            outputs = model.generate(**inputs, max_length=max_length, temperature=temperature)
+            generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
         except Exception as e:
             st.error(f"Error generating text: {str(e)}")
         else:
@@ -41,6 +46,6 @@ if st.button("Generate Text"):
 # About section
 with st.expander("About"):
     st.write("""
-        The Gemma Text Generation app uses the powerful Gemma-7b model from Google to generate text. 
+        This text generation app utilizes the powerful Mistral-8x22B model from the Mistral community on Hugging Face. 
         Adjust the sliders to change the length and creativity of the output.
     """)