Update app.py

app.py

@@ -1,20 +1,47 @@
+import os
 from transformers import pipeline
 from sentence_transformers import SentenceTransformer
-import os
 
-HF_TOKEN = os.getenv("HF_TOKEN")
+# Models
+MAIN_MODEL = "PuruAI/Medini_Intelligence"
+FALLBACK_MODEL = "gpt2"
+EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
 
-# Embedding model
+# Check if token is set
+HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")
+if not HF_TOKEN:
+    print("⚠️ Warning: HUGGINGFACE_HUB_TOKEN not set. Private models may fail to load.")
+
+# --- Load text generation model ---
 try:
-    embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", use_auth_token=HF_TOKEN)
+    generator = pipeline(
+        "text-generation",
+        model=MAIN_MODEL
+    )
+    print(f"✅ Loaded main model: {MAIN_MODEL}")
 except Exception as e:
-    print(f"Failed to load embedding model: {e}")
+    print(f"❌ Failed to load main model: {e}")
+    print(f"⏩ Falling back to {FALLBACK_MODEL}")
+    generator = pipeline("text-generation", model=FALLBACK_MODEL)
+    print(f"✅ Loaded fallback model: {FALLBACK_MODEL}")
 
-# Main model
-MODEL_ID = "PuruAI/Medini_Intelligence"
-FALLBACK_MODEL = "gpt2"
+# --- Load embedding model ---
 try:
-    generator = pipeline("text-generation", model=MODEL_ID, use_auth_token=HF_TOKEN)
+    embedder = SentenceTransformer(EMBEDDING_MODEL)
+    print(f"✅ Loaded embedding model: {EMBEDDING_MODEL}")
 except Exception as e:
-    print(f"Failed to load main model, falling back to GPT-2: {e}")
-    generator = pipeline("text-generation", model=FALLBACK_MODEL)
+    print(f"❌ Failed to load embedding model: {e}")
+    embedder = None  # Safe fallback
+
+# --- Example usage ---
+prompt = "Once upon a time"
+output = generator(prompt, max_length=50)
+print("\n--- Generated Text ---")
+print(output[0]['generated_text'])
+
+if embedder:
+    sentences = ["Hello world", "How are you?"]
+    embeddings = embedder.encode(sentences)
+    print("\n--- Embeddings ---")
+    for s, emb in zip(sentences, embeddings):
+        print(f"{s}: {emb[:5]}...")  # print first 5 values for brevity