Update app.py

app.py

@@ -1,36 +1,34 @@
 import os
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
-from transformers. utils import logging
-
-# Suppress warnings
-logging.set_verbosity_error()
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 # Model configuration
 MODEL_ID = "PuruAI/Medini_Intelligence"
 FALLBACK_MODEL = "gpt2"
 
-# Load token from environment variable
+# Get Hugging Face token from secret/environment variable
 HF_TOKEN = os.environ.get("HUGGINGFACE_HUB_TOKEN")
 
-if HF_TOKEN is None:
-    print("⚠️ Warning: HUGGINGFACE_HUB_TOKEN not set. Private models may fail to load.")
+if not HF_TOKEN:
+    raise ValueError("HUGGINGFACE_HUB_TOKEN is not set. Please set it as a secret/environment variable.")
 
-# Function to load the model safely
 def load_model(model_id):
+    """
+    Try to load the specified model. If it fails, fall back to GPT-2.
+    """
     try:
-        # Load tokenizer and model from Hugging Face
-        tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=HF_TOKEN)
-        model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=HF_TOKEN)
-        print(f"✅ Successfully loaded {model_id}")
+        print(f"Loading model: {model_id}")
+        tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN)
+        model = AutoModelForCausalLM.from_pretrained(model_id, token=HF_TOKEN)
         return pipeline("text-generation", model=model, tokenizer=tokenizer)
     except Exception as e:
         print(f"❌ Failed to load {model_id}: {e}")
         print(f"⏩ Falling back to {FALLBACK_MODEL}")
-        return pipeline("text-generation", model=FALLBACK_MODEL)
+        return pipeline("text-generation", model=FALLBACK_MODEL, token=HF_TOKEN)
 
-# Load Medini AI (with fallback to GPT-2)
+# Initialize the generator
 generator = load_model(MODEL_ID)
 
-# Test the model
-output = generator("Hello, how are you?", max_length=50)
-print(output)
+# Example usage
+prompt = "Once upon a time"
+result = generator(prompt, max_length=50, do_sample=True)
+print(result)