Update app.py

app.py

@@ -1,41 +1,48 @@
+
+import os
 import gradio as gr
 from transformers import pipeline
-import os
 
-# Fetch the API key from environment variables, securely set in Space's settings
+# Load the API key from environment variables
 api_key = os.getenv("HF_API_KEY")
 
-if api_key is not None:
-    os.environ["HF_HOME"] = "/workspace"  # Hugging Face Spaces provides a /workspace directory
+if api_key is None:
+    print("HF_API_KEY not found. Please set it in the Space's secrets.")
+else:
+    # Set the necessary environment variables for Hugging Face Transformers
+    os.environ["HF_HOME"] = "/workspace"  # Use the /workspace directory in Hugging Face Spaces
     os.environ["TRANSFORMERS_CACHE"] = "/workspace/cache"  # Cache directory in Space
-    os.environ["HF_API_KEY"] = api_key
+    os.environ["HF_API_KEY"] = api_key  # Set the API key for Transformers library
 
-    # Attempt to load your model using the pipeline API
-    model_name = "tiiuae/falcon-180B"  # Ensure this is the correct model name
+    # Initialize the model using the API key and Transformers pipeline
     try:
+        model_name = "tiiuae/falcon-180B"  # Ensure this is the correct model name
         nlp_model = pipeline('text-generation', model=model_name)
+        print("Model loaded successfully.")
     except Exception as e:
-        nlp_model = None  # fallback in case of an error
-else:
-    nlp_model = None
-    print("HF_API_KEY not found. Please set it in the Space's secrets.")
+        nlp_model = None
+        print(f"Failed to load model: {str(e)}")
 
 def generate_text(prompt):
+    """Generate text based on the input prompt using the loaded NLP model."""
     if nlp_model:
-        response = nlp_model(prompt, max_length=50, num_return_sequences=1)
-        return response[0]['generated_text']
+        try:
+            response = nlp_model(prompt, max_length=50, num_return_sequences=1)
+            return response[0]['generated_text']
+        except Exception as e:
+            return f"Error during model inference: {str(e)}"
     else:
-        return "Model could not be loaded. Please check the model name and API key."
+        return "Model could not be loaded or initialized properly."
 
-# Gradio interface
+# Setup Gradio interface
 iface = gr.Interface(
     fn=generate_text,
-    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter your text here..."),
+    inputs=gr.Textbox(lines=2, placeholder="Enter your text here..."),
     outputs="text",
     title="Text Generation with Falcon Model",
-    description="Enter some text and see how Falcon model continues it!"
+    description="Enter some text and see how the Falcon model continues it!"
 )
 
-# The following line is essential for Hugging Face Spaces
+# Essential for Hugging Face Spaces
 if __name__ == "__main__":
     iface.launch()