Spaces:
Sleeping
Sleeping
| import os | |
| from openai import OpenAI | |
| # ============================================ | |
| # SETUP — Read environment variables | |
| # ============================================ | |
| API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1") | |
| MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4o-mini") | |
| HF_TOKEN = os.getenv("HF_TOKEN") | |
| if HF_TOKEN is None: | |
| raise ValueError("HF_TOKEN environment variable is required") | |
| # Initialize OpenAI client | |
| client = OpenAI( | |
| base_url=API_BASE_URL, | |
| api_key=HF_TOKEN | |
| ) | |
| # ============================================ | |
| # TASK EMAIL DATASETS | |
| # ============================================ | |
| TASKS = { | |
| "easy_sorting": [ | |
| {"subject": "You won $1,000,000!", "body": "Click here to claim your prize now.", "sender": "prize@randomsite.xyz", "label": "spam"}, | |
| {"subject": "FREE iPhone giveaway", "body": "You have been selected. Claim before midnight!", "sender": "free@giveaway123.com", "label": "spam"}, | |
| {"subject": "Congratulations! You're a winner", "body": "Send your details to collect your reward.", "sender": "win@prizes.net", "label": "spam"}, | |
| {"subject": "Meeting at 3pm today", "body": "Hi, reminder about our team sync at 3pm.", "sender": "manager@company.com", "label": "important"}, | |
| {"subject": "Your invoice is ready", "body": "Please find your monthly invoice attached.", "sender": "billing@service.com", "label": "important"}, | |
| ], | |
| "medium_sorting": [ | |
| {"subject": "50% off this weekend only!", "body": "Flash sale on all items. Use code SAVE50.", "sender": "deals@amazon.com", "label": "promotion"}, | |
| {"subject": "New arrivals just for you", "body": "Check out our latest summer collection.", "sender": "news@shop.com", "label": "promotion"}, | |
| {"subject": "Exclusive member offer inside", "body": "As a valued member, enjoy 20% off.", "sender": "offers@store.com", "label": "promotion"}, | |
| {"subject": "Action required: password expiry", "body": "Your password expires in 3 days. Reset it now.", "sender": "security@company.com", "label": "important"}, | |
| {"subject": "RE: Project update", "body": "Thanks for the update. Let's connect tomorrow.", "sender": "colleague@work.com", "label": "important"}, | |
| {"subject": "Urgent: verify your account", "body": "Your account will be suspended. Click to verify.", "sender": "alert@bank-secure.xyz", "label": "spam"}, | |
| ], | |
| "hard_sorting": [ | |
| {"subject": "Your account statement", "body": "Your monthly statement from XYZ Bank is ready.", "sender": "statements@xyzbank.com", "label": "important"}, | |
| {"subject": "Limited time: upgrade your plan", "body": "Switch to premium and save 30% this month only.", "sender": "offers@service.com", "label": "promotion"}, | |
| {"subject": "Security alert", "body": "A new login was detected from an unknown device.", "sender": "security@google.com", "label": "important"}, | |
| {"subject": "You have unclaimed rewards", "body": "Collect your loyalty points before they expire.", "sender": "rewards@airline.com", "label": "promotion"}, | |
| {"subject": "Final notice: payment overdue", "body": "Send $500 to avoid service interruption.", "sender": "billing@suspicious.xyz", "label": "spam"}, | |
| {"subject": "Team offsite next Friday", "body": "Please confirm your attendance for the offsite.", "sender": "hr@company.com", "label": "important"}, | |
| {"subject": "Claim your free trial", "body": "Start your 30-day free trial — no credit card needed.", "sender": "trial@software.com", "label": "promotion"}, | |
| {"subject": "You've been pre-approved!", "body": "You qualify for a $50,000 loan. Apply now.", "sender": "loans@quickcash.xyz", "label": "spam"}, | |
| ], | |
| } | |
| # ============================================ | |
| # LLM CLASSIFIER | |
| # ============================================ | |
| def classify_email(email: dict) -> str: | |
| prompt = f"""Classify this email into exactly one category: spam, important, or promotion. | |
| Subject: {email['subject']} | |
| From: {email['sender']} | |
| Body: {email['body']} | |
| Reply with ONLY one word: spam, important, or promotion""" | |
| try: | |
| response = client.chat.completions.create( | |
| model=MODEL_NAME, | |
| messages=[ | |
| {"role": "system", "content": "You are an email classifier. Reply with only one word: spam, important, or promotion."}, | |
| {"role": "user", "content": prompt} | |
| ], | |
| max_tokens=10, | |
| temperature=0.0 | |
| ) | |
| answer = response.choices[0].message.content.strip().lower() | |
| if "spam" in answer: return "spam" | |
| if "promotion" in answer: return "promotion" | |
| if "important" in answer: return "important" | |
| return "spam" | |
| except Exception: | |
| return fallback_classify(email) | |
| def fallback_classify(email: dict) -> str: | |
| text = (email["subject"] + " " + email["body"]).lower() | |
| spam_kw = ["won", "free", "prize", "claim", "urgent", "congratulations", | |
| "selected", "suspended", "verify", "overdue", "pre-approved"] | |
| promo_kw = ["off", "sale", "deal", "discount", "offer", "save", | |
| "upgrade", "trial", "rewards", "loyalty"] | |
| spam_score = sum(1 for k in spam_kw if k in text) | |
| promo_score = sum(1 for k in promo_kw if k in text) | |
| if spam_score >= 2: return "spam" | |
| if promo_score >= 1: return "promotion" | |
| return "important" | |
| # ============================================ | |
| # RUN ONE TASK EPISODE | |
| # ============================================ | |
| def run_task(task_id: str, emails: list): | |
| rewards = [] | |
| success = True | |
| # [START] — one per task | |
| print(f"[START] task={task_id} env=email-sorting-openenv model={MODEL_NAME}", flush=True) | |
| try: | |
| for i, email in enumerate(emails, 1): | |
| action = classify_email(email) | |
| correct = (action == email["label"]) | |
| reward = 1.0 if correct else -0.5 | |
| done = (i == len(emails)) | |
| rewards.append(reward) | |
| # [STEP] — one per email | |
| print(f"[STEP] step={i} action={action} reward={reward:.2f} done={'true' if done else 'false'} error=null", flush=True) | |
| except Exception as e: | |
| success = False | |
| correct_count = sum(1 for r in rewards if r > 0) | |
| score = round(correct_count / len(emails), 2) if emails else 0.0 | |
| # Clamp strictly between 0 and 1 | |
| score = round(min(0.99, max(0.01, score)), 2) | |
| rewards_str = ",".join(f"{r:.2f}" for r in rewards) | |
| # [END] — includes score= which validator reads for grader check | |
| print(f"[END] task={task_id} success={'true' if success else 'false'} steps={len(rewards)} score={score} rewards={rewards_str}", flush=True) | |
| return score | |
| # ============================================ | |
| # ENTRY POINT — runs all 3 tasks | |
| # ============================================ | |
| if __name__ == "__main__": | |
| for task_id, emails in TASKS.items(): | |
| run_task(task_id, emails) | |