Upload app.py with huggingface_hub

app.py

@@ -1,51 +1,82 @@
-import os
-import random
-import threading
-
-import discord
 import gradio as gr
-from discord import app_commands
-from discord.ext import commands
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
+tokenizer = AutoTokenizer.from_pretrained("iamhariraj/DialoGPT-small-Rick")
+model = AutoModelForCausalLM.from_pretrained("iamhariraj/DialoGPT-small-Rick")
+model.eval()
 
-# HF GUILD SETTINGS
-MY_GUILD_ID = 1077674588122648679 if os.getenv("TEST_ENV", False) else 879548962464493619
-MY_GUILD = discord.Object(id=MY_GUILD_ID)
-DISCORD_TOKEN = os.environ.get("DISCORD_TOKEN", None)
+def chat(user_message, history):
+    # Build conversation history into input
+    input_ids = None
+    for human, bot in history:
+        h = tokenizer.encode(human + tokenizer.eos_token, return_tensors="pt")
+        b = tokenizer.encode(bot + tokenizer.eos_token, return_tensors="pt")
+        input_ids = torch.cat([input_ids, h, b], dim=-1) if input_ids is not None else torch.cat([h, b], dim=-1)
 
+    new_input = tokenizer.encode(user_message + tokenizer.eos_token, return_tensors="pt")
+    input_ids = torch.cat([input_ids, new_input], dim=-1) if input_ids is not None else new_input
 
-class Bot(commands.Bot):
-    """This structure allows slash commands to work instantly."""
+    # Keep context window manageable
+    if input_ids.shape[-1] > 800:
+        input_ids = input_ids[:, -800:]
 
-    def __init__(self):
-        super().__init__(command_prefix="/", intents=discord.Intents.all())
+    with torch.no_grad():
+        output = model.generate(
+            input_ids,
+            max_length=input_ids.shape[-1] + 100,
+            pad_token_id=tokenizer.eos_token_id,
+            no_repeat_ngram_size=3,
+            do_sample=True,
+            top_k=100,
+            top_p=0.7,
+            temperature=0.8
+        )
 
-    async def setup_hook(self):
-        await self.tree.sync(guild=discord.Object(MY_GUILD_ID))
-        print(f"Synced slash commands for {self.user}.")
+    response = tokenizer.decode(
+        output[:, input_ids.shape[-1]:][0],
+        skip_special_tokens=True
+    )
+    return response
 
 
-client = Bot()
+examples = [
+    "What's the meaning of life?",
+    "Are you smarter than everyone?",
+    "I need your help with something.",
+    "What do you think about Morty?",
+    "Can you build a portal gun?",
+    "What happens when we die?",
+    "Are parallel universes real?",
+]
 
+with gr.Blocks(theme=gr.themes.Monochrome(), title="RickChatBot") as demo:
+    gr.Markdown("""
+    # 🧪 RickChatBot
+    ### Talk to an AI Rick Sanchez — *the smartest being in the universe*
+    > Fine-tuned on Rick & Morty dialogue. Powered by DialoGPT.
+    """)
 
-@client.event
-async def on_ready():
-    print(f"Logged in as {client.user} (ID: {client.user.id})")
-    print("------")
+    chatbot = gr.Chatbot(height=400, label="Rick Sanchez")
+    msg = gr.Textbox(placeholder="Say something to Rick...", label="You", lines=1)
 
+    with gr.Row():
+        send = gr.Button("Send", variant="primary")
+        clear = gr.Button("Clear Chat")
 
-def run_bot():
-    client.run(DISCORD_TOKEN)
+    gr.Examples(examples=examples, inputs=msg)
 
+    gr.Markdown("> ⚠️ AI-generated responses. May take ~30s on first load (Space waking up).")
 
-threading.Thread(target=run_bot).start()
-"""This allows us to run the Discord bot in a Python thread"""
-with gr.Blocks() as demo:
-    gr.Markdown("""
-    # Huggingbots Server
-    This space hosts the huggingbots discord bot.
-    Currently supported models are Falcon and DeepfloydIF
-    """)
-demo.queue(concurrency_count=100)
-demo.queue(max_size=100)
-demo.launch()
+    def respond(message, chat_history):
+        if not message.strip():
+            return "", chat_history
+        bot_response = chat(message, chat_history)
+        chat_history.append((message, bot_response))
+        return "", chat_history
+
+    send.click(respond, [msg, chatbot], [msg, chatbot])
+    msg.submit(respond, [msg, chatbot], [msg, chatbot])
+    clear.click(lambda: [], None, chatbot)
+
+demo.launch()