Upload discord_bot.py

discord_bot.py

@@ -0,0 +1,224 @@
+# -*- coding: utf-8 -*-
+"""
+Text2Receipt — Discord Bot  (+5% bonus)
+========================================
+Run this cell in Google Colab (T4) AFTER the fine-tuned model is loaded.
+The bot wraps the same pipeline as the Gradio Space.
+
+Usage in any Discord channel:
+  !receipt <Hebrew income note>
+
+Example:
+  !receipt קיבלתי 1200 שקל ממשה כהן על ייעוץ עסקי
+
+Setup:
+  1. Go to https://discord.com/developers/applications → New Application
+  2. Bot tab → Add Bot → copy token
+  3. OAuth2 → URL Generator: scope=bot, permissions=Send Messages
+  4. Paste DISCORD_TOKEN below (or set as Colab env var)
+  5. Run this cell while the model is already loaded in the session
+"""
+
+import os, json, random, datetime as _dt, asyncio, threading
+import discord
+from discord.ext import commands
+
+# ── same pipeline helpers as app.py ──────────────────────────────────────────
+import t2r_core as core
+
+DEMO_ISSUER = {
+    "name":   "מים שקטים",
+    "tax_id": "962569844",
+    "status": "authorized_dealer",
+}
+
+_RNG = random.Random(99)
+
+# These are defined here so the bot works standalone (copied from app.py)
+INSTRUCTION = (
+    "אתה ממיר הערת הכנסה חופשית בעברית למבנה JSON. "
+    "חלץ אך ורק את מה שכתוב בהערה: שם הלקוח (client_name), "
+    "האם הלקוח עסק (client_is_business), ורשימת פריטים (items) "
+    "כאשר לכל פריט תיאור (description), מחיר ליחידה (unit_price) וכמות (quantity). "
+    "החזר JSON תקין בלבד, ללא טקסט נוסף."
+)
+
+def _build_prompt(raw_text: str) -> str:
+    return f"{INSTRUCTION}\n\nהערה: {raw_text}\n\nJSON:"
+
+def _extract_json(text: str):
+    s = text.find("{")
+    if s < 0:
+        return None
+    depth = 0
+    for i in range(s, len(text)):
+        if text[i] == "{":
+            depth += 1
+        elif text[i] == "}":
+            depth -= 1
+            if depth == 0:
+                try:
+                    return json.loads(text[s:i+1])
+                except Exception:
+                    return None
+    return None
+
+def _apply_defaults(parse: dict) -> dict:
+    p = dict(parse)
+    p.setdefault("doc_type",           "receipt")
+    p.setdefault("date",               _dt.date.today().isoformat())
+    p.setdefault("payment_method",     "bank_transfer")
+    p.setdefault("amount_basis",       "net")
+    p.setdefault("currency",           "ILS")
+    p.setdefault("client_tax_id",      None)
+    p.setdefault("client_is_business", False)
+    p.setdefault("items",              [])
+    return p
+
+def _format_document(completed: dict) -> str:
+    """Format completed document as Discord-friendly Markdown."""
+    c       = completed
+    issuer  = c.get("issuer", {})
+    client  = c.get("client", {})
+    lines   = c.get("lines",  [])
+    vat_pct = int(round(c.get("vat_rate", 0) * 100))
+
+    line_rows = "\n".join(
+        f"  • {ln['description']} × {ln['quantity']} → ₪{ln['line_total']:,.2f}"
+        for ln in lines
+    )
+    alloc = (
+        f"\n🔖 **מספר הקצאה:** `{c.get('allocation_number')}`"
+        if c.get("allocation_required") else ""
+    )
+
+    return (
+        f"```\n"
+        f"{'━'*38}\n"
+        f"  {c.get('doc_type_he','מסמך')}   |   מס׳ {c.get('serial_number','')}\n"
+        f"  {c.get('issue_date','')}\n"
+        f"{'━'*38}\n"
+        f"  מנפיק : {issuer.get('name','')}  (ח.פ. {issuer.get('tax_id','')})\n"
+        f"  לקוח   : {client.get('name','')}"
+        f"{'  [עסק]' if client.get('is_business') else ''}\n"
+        f"{'━'*38}\n"
+        f"{line_rows}\n"
+        f"{'━'*38}\n"
+        f"  סכום נטו  : ₪{c.get('subtotal',0):>10,.2f}\n"
+        f"  מע\"מ {vat_pct:2d}%  : ₪{c.get('vat_amount',0):>10,.2f}\n"
+        f"  **סה\"כ**   : ₪{c.get('total',0):>10,.2f}\n"
+        f"  תשלום : {core.PAYMENT_HE.get(c.get('payment_method',''),'—')}\n"
+        f"{'━'*38}\n"
+        f"```"
+        f"{alloc}"
+    )
+
+
+def run_pipeline(raw_text: str, tok, model) -> str:
+    """
+    Full pipeline: raw Hebrew note → formatted fiscal document string.
+    tok and model are the already-loaded objects from the Colab session.
+    """
+    import torch
+
+    # 1. Parse
+    msgs   = [{"role": "user", "content": _build_prompt(raw_text)}]
+    prompt = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+    enc    = tok(prompt, return_tensors="pt",
+                 add_special_tokens=False).to(model.device)
+    with torch.no_grad():
+        out = model.generate(**enc, max_new_tokens=200, do_sample=False,
+                             pad_token_id=tok.pad_token_id)
+    decoded = tok.decode(out[0, enc["input_ids"].shape[1]:],
+                         skip_special_tokens=True)
+    parse = _extract_json(decoded)
+
+    if parse is None:
+        return "❌ לא הצלחתי לחלץ את הפרטים. נסה לנסח מחדש."
+
+    # 2. Apply defaults + complete
+    final_parse = _apply_defaults(parse)
+    try:
+        completed = core.complete(DEMO_ISSUER, final_parse, _RNG)
+    except Exception as e:
+        return f"❌ שגיאה בעיבוד: {e}"
+
+    # 3. Format
+    return _format_document(completed)
+
+
+def start_bot(tok, model, discord_token: str = None):
+    """
+    Start the Discord bot. Call this from Colab after loading tok + model.
+
+    Args:
+        tok:           HuggingFace tokenizer (already loaded)
+        model:         fine-tuned model (already loaded, on GPU)
+        discord_token: Bot token. Falls back to DISCORD_TOKEN env var.
+    """
+    token = discord_token or os.environ.get("DISCORD_TOKEN", "")
+    if not token:
+        raise ValueError(
+            "No Discord token found. "
+            "Pass discord_token= or set os.environ['DISCORD_TOKEN']"
+        )
+
+    intents = discord.Intents.default()
+    intents.message_content = True
+    bot = commands.Bot(command_prefix="!", intents=intents)
+
+    @bot.event
+    async def on_ready():
+        print(f"✅ Discord bot ready: {bot.user} (id: {bot.user.id})")
+        print("   Listening for: !receipt <Hebrew note>")
+
+    @bot.command(name="receipt")
+    async def receipt_cmd(ctx, *, note: str):
+        """Convert a Hebrew income note to a fiscal document."""
+        await ctx.message.add_reaction("⏳")
+        try:
+            # run blocking inference in executor to not block event loop
+            loop   = asyncio.get_event_loop()
+            result = await loop.run_in_executor(
+                None, run_pipeline, note, tok, model
+            )
+            await ctx.reply(result)
+        except Exception as e:
+            await ctx.reply(f"❌ שגיאה: {e}")
+        finally:
+            await ctx.message.remove_reaction("⏳", bot.user)
+            await ctx.message.add_reaction("✅")
+
+    @bot.command(name="help_t2r")
+    async def help_cmd(ctx):
+        await ctx.reply(
+            "**Text2Receipt Bot** 🧾\n"
+            "המר הערת הכנסה עברית למסמך פיסקלי ישראלי.\n\n"
+            "שימוש: `!receipt <הערה בעברית>`\n\n"
+            "דוגמאות:\n"
+            "• `!receipt קיבלתי 500 שקל ממשה על ייעוץ`\n"
+            "• `!receipt מעגל בע\"מ שילמה 15,000 ש\"ח על פרויקט אתר`"
+        )
+
+    # Run in background thread so Colab cell doesn't block
+    def _run():
+        asyncio.run(bot.start(token))
+
+    t = threading.Thread(target=_run, daemon=True)
+    t.start()
+    print("🤖 Bot thread started. Keep this cell running.")
+    return bot
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# Colab usage example (paste into a cell after loading tok + model in nb03):
+# ═════════════════════════════════════════════════════════════════════════════
+# import os
+# os.environ["DISCORD_TOKEN"] = "YOUR_BOT_TOKEN_HERE"
+#
+# from discord_bot import start_bot
+# bot = start_bot(tok, ft_model)   # tok + ft_model from §5 of nb03
+#
+# # Test locally without Discord:
+# from discord_bot import run_pipeline
+# print(run_pipeline("קיבלתי 1200 שקל ממשה כהן על ייעוץ עסקי", tok, ft_model))