Update app.py

app.py

@@ -1,34 +1,29 @@
 import os
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 
 # Model configuration
-MODEL_ID = "PuruAI/Medini_Intelligence"
+MODEL_ID = "PuruAI/Medini_Intelligence"  # Your private model
 FALLBACK_MODEL = "gpt2"
 
-# Get Hugging Face token from secret/environment variable
-HF_TOKEN = os.environ.get("HUGGINGFACE_HUB_TOKEN")
-
-if not HF_TOKEN:
-    raise ValueError("HUGGINGFACE_HUB_TOKEN is not set. Please set it as a secret/environment variable.")
+# Load Hugging Face token from secret
+HF_TOKEN = os.environ.get("HF_TOKEN")
 
 def load_model(model_id):
-    """
-    Try to load the specified model. If it fails, fall back to GPT-2.
-    """
+    """Load private model if possible, otherwise fallback to GPT-2."""
     try:
         print(f"Loading model: {model_id}")
-        tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN)
-        model = AutoModelForCausalLM.from_pretrained(model_id, token=HF_TOKEN)
+        tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=HF_TOKEN)
+        model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_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, token=HF_TOKEN)
+        return pipeline("text-generation", model=FALLBACK_MODEL)
 
-# Initialize the generator
+# Initialize pipeline
 generator = load_model(MODEL_ID)
 
-# Example usage
+# Test generation
 prompt = "Once upon a time"
-result = generator(prompt, max_length=50, do_sample=True)
-print(result)
+output = generator(prompt, max_new_tokens=50)
+print(output)