app.py

import os
import gradio as gr
from transformers import pipeline

# Load the API key from environment variables
api_key = os.getenv("HF_API_KEY")

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  # Set the API key for Transformers library

    # 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
        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:
        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 or initialized properly."

# Setup Gradio interface
iface = gr.Interface(
    fn=generate_text,
    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 the Falcon model continues it!"
)

# Essential for Hugging Face Spaces
if __name__ == "__main__":
    iface.launch()