Update app.py

app.py

@@ -1,53 +1,49 @@
+import os
 import discord
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from huggingface_hub import login
+from dotenv import load_dotenv
 
-# Hugging Face Token (replace 'your-huggingface-token' with your actual token)
-HF_TOKEN = "your-huggingface-token"
+# Load environment variables
+load_dotenv()
+HF_TOKEN = os.getenv("HF_TOKEN")
+DISCORD_TOKEN = os.getenv("DISCORD_TOKEN")
+
+if not HF_TOKEN:
+    raise ValueError("Hugging Face token is missing. Set HF_TOKEN in the environment variables.")
+
+if not DISCORD_TOKEN:
+    raise ValueError("Discord bot token is missing. Set DISCORD_TOKEN in the environment variables.")
+
+# Authenticate with Hugging Face
 login(HF_TOKEN)
 
-# Load DeepScaleR model from Hugging Face
-MODEL_NAME = "DeepScale/DeepScaleR"
+# Load DeepScaleR model and tokenizer
+MODEL_NAME = "Your_HuggingFace_Repo/DeepScaleR"  # Replace with your model repo name
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16)
-
-# Discord Bot Token (replace 'your-discord-token' with your actual bot token)
-DISCORD_TOKEN = "your-discord-token"
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16).to("cuda" if torch.cuda.is_available() else "cpu")
 
-# Set up Discord bot
+# Initialize Discord bot
 intents = discord.Intents.default()
 intents.messages = True
 client = discord.Client(intents=intents)
 
-# Response rules for Shiv Yantra AI
-async def respond(message):
-    if message.author == client.user:
-        return  # Ignore itself
-
-    user_input = message.content.strip()
-    inputs = tokenizer(user_input, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
-    outputs = model.generate(**inputs, max_length=500)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Modify response to match Shiv Yantra AI's identity
-    if "DeepScaleR" in response:
-        response = response.replace("DeepScaleR", "Shiv Yantra AI")
-    if "Who made you?" in user_input:
-        response = "I was created by Spectral Satya."
-    if "Who is your founder?" in user_input:
-        response = "My founder is Hardik Kumawat."
-
-    await message.channel.send(response)
-
-# Discord event handlers
 @client.event
 async def on_ready():
-    print(f"Logged in as {client.user}")
+    print(f'Logged in as {client.user}')
 
 @client.event
 async def on_message(message):
-    await respond(message)
+    if message.author == client.user:
+        return
+
+    input_text = message.content
+    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+    output = model.generate(**inputs, max_length=200)
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+
+    await message.channel.send(response)
 
-# Start bot
+# Run Discord bot
 client.run(DISCORD_TOKEN)