app.py

import os
import threading
import discord
import requests
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from dotenv import load_dotenv

# 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 model repo is private

if not DISCORD_TOKEN:
    raise ValueError("Discord bot token is missing. Set DISCORD_TOKEN in the environment variables.")

# Set the model repository name. For DeepScaleR-1.5B-Preview, use:
MODEL_NAME = "agentica-org/DeepScaleR-1.5B-Preview"

# 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, 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"
    )

# 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)
    # Replace any instance of the internal model name with your bot's identity.
    response = response.replace("DeepScaleR", "Shiv Yantra AI")
    return response

# ==========================
# Gradio API Setup
# ==========================
def gradio_api(input_text):
    return generate_response(input_text)

iface = gr.Interface(fn=gradio_api, inputs="text", outputs="text")

def run_gradio():
    iface.launch(server_name="0.0.0.0", server_port=7860, share=False)

# ==========================
# Discord Bot Setup
# ==========================
intents = discord.Intents.default()
intents.message_content = True  # Needed to read message contents
client = discord.Client(intents=intents)

# Local endpoint for the Gradio API
GRADIO_API_URL = "http://0.0.0.0:7860/run/predict"

@client.event
async def on_ready():
    print(f"Logged in as {client.user}")

@client.event
async def on_message(message):
    if message.author == client.user:
        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 communicating with the AI API."
        await message.channel.send(ai_response)

def run_discord_bot():
    client.run(DISCORD_TOKEN)

# ==========================
# Start Both Services
# ==========================
if __name__ == "__main__":
    # Start the Gradio interface in a daemon thread
    threading.Thread(target=run_gradio, daemon=True).start()
    # Start the Discord bot in a daemon thread
    threading.Thread(target=run_discord_bot, daemon=True).start()

    # Keep the main thread alive.
    while True:
        pass