Update app.py

app.py

@@ -1,23 +1,36 @@
 import os
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from transformers.utils import logging
 
-# ✅ Set your Hugging Face token before any Hugging Face imports
-os.environ["HUGGINGFACE_HUB_TOKEN"] = "YOUR_HUGGINGFACE_TOKEN"  # Replace with your actual token
-
-from transformers import pipeline
+# Suppress warnings
+logging.set_verbosity_error()
 
 # Model configuration
 MODEL_ID = "PuruAI/Medini_Intelligence"
 FALLBACK_MODEL = "gpt2"
 
-# Try loading your main model
-try:
-    generator = pipeline("text-generation", model=MODEL_ID)
-    print("✅ Medini_AI loaded successfully!")
-except Exception as e:
-    print(f"⚠️ Failed to load Medini_AI: {e}")
-    print("⏩ Falling back to GPT-2...")
-    generator = pipeline("text-generation", model=FALLBACK_MODEL)
+# Load token from 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.")
+
+# Function to load the model safely
+def load_model(model_id):
+    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}")
+        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)
+
+# Load Medini AI (with fallback to GPT-2)
+generator = load_model(MODEL_ID)
 
-# Test the pipeline
-result = generator("Hello world!", max_length=50)
-print(result)
+# Test the model
+output = generator("Hello, how are you?", max_length=50)
+print(output)