Spaces:
Running
Running
| import os | |
| import threading | |
| import discord | |
| 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 available) | |
| load_dotenv() | |
| DISCORD_TOKEN = os.getenv("DISCORD_TOKEN") | |
| HF_TOKEN = os.getenv("HF_TOKEN") # Optional: only if needed for private models | |
| if not DISCORD_TOKEN: | |
| raise ValueError("Discord bot token is missing. Set DISCORD_TOKEN in the environment variables.") | |
| # Set the model repository name (public model) | |
| MODEL_NAME = "agentica-org/DeepScaleR-1.5B-Preview" | |
| # Load the tokenizer and model. Use token=HF_TOKEN if provided. | |
| 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" | |
| ) | |
| # Function to generate AI responses | |
| 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 always identifies as "Shiv Yantra AI" | |
| response = response.replace("DeepScaleR", "Shiv Yantra AI") | |
| return response | |
| # ========================== | |
| # Gradio Interface (optional UI) | |
| # ========================== | |
| def gradio_api(input_text): | |
| return generate_response(input_text) | |
| iface = gr.Interface(fn=gradio_api, inputs="text", outputs="text", title="Shiv Yantra AI") | |
| def run_gradio(): | |
| iface.launch(server_name="0.0.0.0", server_port=7860, share=False) | |
| # ========================== | |
| # Discord Bot Setup (Directly uses local generate_response) | |
| # ========================== | |
| intents = discord.Intents.default() | |
| intents.message_content = True # Required for reading message content | |
| client = discord.Client(intents=intents) | |
| async def on_ready(): | |
| print(f"Logged in as {client.user}") | |
| async def on_message(message): | |
| # Avoid replying to itself | |
| if message.author == client.user: | |
| return | |
| user_input = message.content.strip() | |
| if user_input: | |
| try: | |
| # Directly call the local generate_response function | |
| ai_response = generate_response(user_input) | |
| except Exception as e: | |
| ai_response = "Error processing your request." | |
| await message.channel.send(ai_response) | |
| def run_discord_bot(): | |
| client.run(DISCORD_TOKEN) | |
| # ========================== | |
| # Start Both Services Concurrently | |
| # ========================== | |
| if __name__ == "__main__": | |
| # Start the Gradio interface in a separate thread (optional UI) | |
| threading.Thread(target=run_gradio, daemon=True).start() | |
| # Start the Discord bot in a separate thread | |
| threading.Thread(target=run_discord_bot, daemon=True).start() | |
| # Keep the main thread alive indefinitely | |
| while True: | |
| pass |