Update app.py

app.py

@@ -1,49 +1,93 @@
 import os
+import threading
 import discord
+import requests
 import torch
+import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from huggingface_hub import login
-from dotenv import load_dotenv
 
-# Load environment variables
-load_dotenv()
-HF_TOKEN = os.getenv("HF_TOKEN")
+# Load tokens from environment variables
 DISCORD_TOKEN = os.getenv("DISCORD_TOKEN")
-
-if not HF_TOKEN:
-    raise ValueError("Hugging Face token is missing. Set HF_TOKEN in the environment variables.")
+HF_TOKEN = os.getenv("HF_TOKEN")  # Optional: only needed if your model is private
 
 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)
+# Use the official DeepScaleR model repository
+MODEL_NAME = "agentica-org/DeepScaleR-1.5B-Preview"
+
+# Load the model and tokenizer (if HF_TOKEN is provided, use it for authentication)
+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")
+else:
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, device_map="auto")
+
+# Define a function to generate a response using 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"
+    response = response.replace("DeepScaleR", "Shiv Yantra AI")
+    return response
+
+# ==========================
+# Gradio API Setup
+# ==========================
+def gradio_api(input_text):
+    return generate_response(input_text)
 
-# 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).to("cuda" if torch.cuda.is_available() else "cpu")
+iface = gr.Interface(fn=gradio_api, inputs="text", outputs="text")
 
-# Initialize Discord bot
+def run_gradio():
+    iface.launch(server_name="0.0.0.0", server_port=7860, share=False)
+
+# ==========================
+# Discord Bot Setup
+# ==========================
 intents = discord.Intents.default()
-intents.messages = True
+intents.message_content = True  # Required for reading message content
 client = discord.Client(intents=intents)
 
+# Use the local Gradio API endpoint (since both run in the same Space)
+GRADIO_API_URL = "http://0.0.0.0:7860/run/predict"
+
 @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):
     if message.author == client.user:
-        return
+        return  # Avoid replying to self
+
+    user_input = message.content.strip()
+    if user_input:
+        try:
+            payload = {"data": [user_input]}
+            r = requests.post(GRADIO_API_URL, json=payload)
+            r.raise_for_status()
+            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."
+        await message.channel.send(ai_response)
 
-    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)
+def run_discord_bot():
+    client.run(DISCORD_TOKEN)
 
-    await message.channel.send(response)
+# ==========================
+# Start Both Services
+# ==========================
+if __name__ == "__main__":
+    # Start Gradio in a separate daemon thread
+    threading.Thread(target=run_gradio, daemon=True).start()
+    # Start Discord bot in another daemon thread
+    threading.Thread(target=run_discord_bot, daemon=True).start()
 
-# Run Discord bot
-client.run(DISCORD_TOKEN)
+    # Keep the main thread alive indefinitely
+    while True:
+        pass