Update app.py

app.py

@@ -5,32 +5,39 @@ import requests
 import torch
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from dotenv import load_dotenv
 
-# Load tokens from environment variables
+# Load environment variables from Hugging Face Secrets (and .env if local)
+load_dotenv()
 DISCORD_TOKEN = os.getenv("DISCORD_TOKEN")
-HF_TOKEN = os.getenv("HF_TOKEN")  # Optional: only needed if your model is private
+HF_TOKEN = os.getenv("HF_TOKEN")  # Optional: only needed if model repo is private
 
 if not DISCORD_TOKEN:
     raise ValueError("Discord bot token is missing. Set DISCORD_TOKEN in the environment variables.")
 
-# Use the official DeepScaleR model repository
+# Set the model repository name. For DeepScaleR-1.5B-Preview, use:
 MODEL_NAME = "agentica-org/DeepScaleR-1.5B-Preview"
 
-# Load the model and tokenizer (if HF_TOKEN is provided, use it for authentication)
+# Load the tokenizer and model.
+# Using token=HF_TOKEN instead of use_auth_token (per the new deprecation)
 if HF_TOKEN:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=HF_TOKEN)
-    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, use_auth_token=HF_TOKEN, torch_dtype=torch.float16, device_map="auto")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME, token=HF_TOKEN, torch_dtype=torch.float16, device_map="auto"
+    )
 else:
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, device_map="auto")
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME, torch_dtype=torch.float16, device_map="auto"
+    )
 
-# Define a function to generate a response using the model
+# Define a function to generate responses with the model
 def generate_response(prompt):
     device = "cuda" if torch.cuda.is_available() else "cpu"
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024).to(device)
     outputs = model.generate(**inputs, max_new_tokens=200, do_sample=True, top_p=0.9, temperature=0.7)
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    # Ensure the bot identifies as "Shiv Yantra AI"
+    # Replace any instance of the internal model name with your bot's identity.
     response = response.replace("DeepScaleR", "Shiv Yantra AI")
     return response
 
@@ -49,10 +56,10 @@ def run_gradio():
 # Discord Bot Setup
 # ==========================
 intents = discord.Intents.default()
-intents.message_content = True  # Required for reading message content
+intents.message_content = True  # Needed to read message contents
 client = discord.Client(intents=intents)
 
-# Use the local Gradio API endpoint (since both run in the same Space)
+# Local endpoint for the Gradio API
 GRADIO_API_URL = "http://0.0.0.0:7860/run/predict"
 
 @client.event
@@ -73,7 +80,7 @@ async def on_message(message):
             response_json = r.json()
             ai_response = response_json.get("data", ["Sorry, something went wrong."])[0]
         except Exception as e:
-            ai_response = "Error contacting the AI API."
+            ai_response = "Error communicating with the AI API."
         await message.channel.send(ai_response)
 
 def run_discord_bot():
@@ -83,11 +90,11 @@ def run_discord_bot():
 # Start Both Services
 # ==========================
 if __name__ == "__main__":
-    # Start Gradio in a separate daemon thread
+    # Start the Gradio interface in a daemon thread
     threading.Thread(target=run_gradio, daemon=True).start()
-    # Start Discord bot in another daemon thread
+    # Start the Discord bot in a daemon thread
     threading.Thread(target=run_discord_bot, daemon=True).start()
 
-    # Keep the main thread alive indefinitely
+    # Keep the main thread alive.
     while True:
         pass