Update app.py

app.py

@@ -1,50 +1,20 @@
-import gradio as gr
 from transformers import pipeline
 from sentence_transformers import SentenceTransformer
-import logging
+import os
 
-# Logging setup
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-# Model configuration
-MODEL_ID = "PuruAI/Medini_Intelligence"
-FALLBACK_MODEL = "gpt2"
-EMB_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
-
-# Load embedding model
+# Embedding model
 try:
-    logger.info(f"Loading embedding model: {EMB_MODEL}")
-    embedding_model = SentenceTransformer(EMB_MODEL)
-    logger.info("Embedding model loaded successfully.")
+    embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", use_auth_token=HF_TOKEN)
 except Exception as e:
-    logger.error(f"Failed to load embedding model: {e}")
-    embedding_model = None
+    print(f"Failed to load embedding model: {e}")
 
-# Load main LLM
+# Main model
+MODEL_ID = "PuruAI/Medini_Intelligence"
+FALLBACK_MODEL = "gpt2"
 try:
-    logger.info(f"Loading main model: {MODEL_ID}")
-    generator = pipeline("text-generation", model=MODEL_ID)
-    logger.info("Main model loaded successfully.")
+    generator = pipeline("text-generation", model=MODEL_ID, use_auth_token=HF_TOKEN)
 except Exception as e:
-    logger.warning(f"Failed to load {MODEL_ID}, falling back to {FALLBACK_MODEL}: {e}")
+    print(f"Failed to load main model, falling back to GPT-2: {e}")
     generator = pipeline("text-generation", model=FALLBACK_MODEL)
-
-# Gradio interface
-def generate_text(prompt):
-    try:
-        result = generator(prompt, max_length=200)
-        return result[0]["generated_text"]
-    except Exception as e:
-        logger.error(f"Text generation failed: {e}")
-        return "Error generating text."
-
-with gr.Blocks() as demo:
-    gr.Markdown("## Medini AI Space")
-    user_input = gr.Textbox(label="Enter your prompt")
-    output = gr.Textbox(label="Generated Text")
-    submit = gr.Button("Generate")
-    submit.click(generate_text, inputs=user_input, outputs=output)
-
-if __name__ == "__main__":
-    demo.launch()