Update app.py

app.py

@@ -1,23 +1,41 @@
-# app.py
-
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline
+import os
+
+# Fetch the API key from environment variables, securely set in Space's settings
+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
+    os.environ["TRANSFORMERS_CACHE"] = "/workspace/cache"  # Cache directory in Space
+    os.environ["HF_API_KEY"] = api_key
 
-def main():
-    # Load the model and tokenizer
-    model_name = "/tiiuae/falcon-180B"  # Replace with the correct model name if different
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
+    # Attempt to load your model using the pipeline API
+    model_name = "tiiuae/falcon-180B"  # Ensure this is the correct model name
+    try:
+        nlp_model = pipeline('text-generation', model=model_name)
+    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.")
 
-    # Define the function that uses the model
-    def generate_text(input_text):
-        inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
-        outputs = model.generate(**inputs, max_length=512, min_length=50, num_return_sequences=1)
-        return tokenizer.decode(outputs[0], skip_special_tokens=True)
+def generate_text(prompt):
+    if nlp_model:
+        response = nlp_model(prompt, max_length=50, num_return_sequences=1)
+        return response[0]['generated_text']
+    else:
+        return "Model could not be loaded. Please check the model name and API key."
 
-    # Create and launch the Gradio interface
-    iface = gr.Interface(fn=generate_text, inputs="text", outputs="text", title="Falcon 180B Text Generator")
-    iface.launch()
+# Gradio interface
+iface = gr.Interface(
+    fn=generate_text,
+    inputs=gr.inputs.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!"
+)
 
+# The following line is essential for Hugging Face Spaces
 if __name__ == "__main__":
-    main()
+    iface.launch()