Update app.py

app.py

@@ -1,44 +1,34 @@
 import os
 from transformers import pipeline
-from sentence_transformers import SentenceTransformer
 
-# --- Hugging Face Models ---
-Medini_AI = "PuruAI/Medini_Intelligence"       # Main model (variable name cleaned)
-FALLBACK_MODEL = "gpt2"                        # Fallback model
-EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+# ✅ Set Hugging Face token (replace with your token)
+# Option 1: Directly in script (works locally)
+os.environ["HUGGINGFACE_HUB_TOKEN"] = "YOUR_HUGGINGFACE_TOKEN_HERE"
 
-# --- Hugging Face Token from Environment ---
-HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")
-if not HF_TOKEN:
-    print("⚠️ Warning: HUGGINGFACE_HUB_TOKEN not set. Private models may fail to load.")
+# Option 2: If running in a container, make sure HUGGINGFACE_HUB_TOKEN
+# is set as an environment variable externally (Docker, cloud, etc.)
 
-# --- Load Main Text Generation Model ---
-try:
-    generator = pipeline("text-generation", model=Medini_AI)
-    print(f"✅ Loaded main model: Medini_AI ({Medini_AI})")
-except Exception as e:
-    print(f"❌ Failed to load Medini_AI: {e}")
-    print(f"⏩ Falling back to {FALLBACK_MODEL}")
-    generator = pipeline("text-generation", model=FALLBACK_MODEL)
-    print(f"✅ Loaded fallback model: {FALLBACK_MODEL}")
+# Model names
+MEDINI_MODEL = "PuruAI/Medini_Intelligence"
+FALLBACK_MODEL = "gpt2"
 
-# --- Load Embedding Model ---
-try:
-    embedder = SentenceTransformer(EMBEDDING_MODEL)
-    print(f"✅ Loaded embedding model: {EMBEDDING_MODEL}")
-except Exception as e:
-    print(f"❌ Failed to load embedding model: {e}")
-    embedder = None  # Safe fallback
+# Function to safely load a model
+def load_model(model_name, fallback_model):
+    try:
+        print(f"Loading model: {model_name}")
+        generator = pipeline("text-generation", model=model_name)
+        print("✅ Model loaded successfully!")
+        return generator
+    except Exception as e:
+        print(f"⚠️ Failed to load {model_name}: {e}")
+        print(f"⏩ Falling back to {fallback_model}")
+        generator = pipeline("text-generation", model=fallback_model)
+        return generator
 
-# --- Example Usage ---
-prompt = "Once upon a time"
-output = generator(prompt, max_length=50)
-print("\n--- Generated Text ---")
-print(output[0]['generated_text'])
+# Load Medini model or fallback
+generator = load_model(MEDINI_MODEL, FALLBACK_MODEL)
 
-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
+# Example usage
+prompt = "Hello Medini AI, can you explain AI agents?"
+result = generator(prompt, max_length=100, do_sample=True)
+print(result)