Upload qwen_anonymizer_fixed.py

qwen_anonymizer_fixed.py

@@ -0,0 +1,341 @@
+"""
+🤖 Qwen Anonymizer Module - نسخه Space (لود لوکال)
+ماژول استفاده مستقیم از مدل فاین‌تیون شده Qwen2.5-1.5B
+✅ بدون نیاز به Inference API - مدل مستقیماً از Space لود می‌شود
+"""
+
+import os
+import logging
+from typing import Optional, Dict, Tuple
+import json
+import re
+import torch
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class QwenAnonymizer:
+    """کلاس برای استفاده از مدل فاین‌تیون شده Qwen - لود لوکال"""
+    
+    def __init__(
+        self,
+        model_path: str = "./qwen-anonymizer-v2",  # ⭐ مسیر لوکال در Space
+        device: str = "auto"
+    ):
+        """
+        مقداردهی اولیه
+        
+        Args:
+            model_path: مسیر مدل در Space
+            device: دستگاه (auto, cuda, cpu)
+        """
+        self.model_path = model_path
+        self.device = device
+        
+        self.model = None
+        self.tokenizer = None
+        self.base_model = None
+        self._model_loaded = False
+        
+        # لود مدل
+        self._load_model()
+        
+        logger.info(f"✅ QwenAnonymizer آماده است: {self.model_path}")
+    
+    def _load_model(self):
+        """لود کردن مدل LoRA + Base Model"""
+        try:
+            from transformers import AutoTokenizer, AutoModelForCausalLM
+            from peft import PeftModel
+            
+            logger.info("📥 شروع لود مدل...")
+            logger.info(f"   مسیر مدل: {self.model_path}")
+            
+            # ✅ مرحله 1: لود Base Model
+            logger.info("📥 لود Base Model: Qwen/Qwen2.5-1.5B")
+            self.base_model = AutoModelForCausalLM.from_pretrained(
+                "Qwen/Qwen2.5-1.5B",
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                device_map=self.device,
+                trust_remote_code=True,
+                low_cpu_mem_usage=True  # کاهش مصرف RAM
+            )
+            logger.info("✅ Base Model لود شد")
+            
+            # ✅ مرحله 2: لود LoRA Adapter
+            logger.info(f"🔧 لود LoRA Adapter از: {self.model_path}")
+            self.model = PeftModel.from_pretrained(
+                self.base_model,
+                self.model_path,
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+            )
+            logger.info("✅ LoRA Adapter لود شد")
+            
+            # ✅ مرحله 3: لود Tokenizer
+            logger.info("📝 لود Tokenizer...")
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                "Qwen/Qwen2.5-1.5B",
+                trust_remote_code=True,
+                fix_mistral_regex=True  # ⭐ رفع مشکل tokenizer
+            )
+            logger.info("✅ Tokenizer لود شد")
+            
+            # تنظیم pad token
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+            
+            # فعال کردن eval mode
+            self.model.eval()
+            
+            self._model_loaded = True
+            
+            # نمایش اطلاعات دستگاه
+            if torch.cuda.is_available():
+                logger.info(f"🚀 استفاده از GPU: {torch.cuda.get_device_name(0)}")
+                logger.info(f"💾 VRAM موجود: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
+            else:
+                logger.info("⚠️  استفاده از CPU (ممکنه کند باشه)")
+            
+            logger.info("✅ مدل کاملاً آماده است!")
+            
+        except Exception as e:
+            logger.error(f"❌ خطا در لود مدل: {e}")
+            logger.error("💡 مطمئن شوید:")
+            logger.error(f"   1. مدل در مسیر {self.model_path} وجود دارد")
+            logger.error("   2. فایل‌های adapter_config.json و adapter_model.safetensors موجود هستند")
+            logger.error("   3. GPU در دسترس است (برای Space)")
+            self._model_loaded = False
+            raise
+    
+    def anonymize(
+        self,
+        text: str,
+        entity_types: list = None,
+        entities_to_anonymize: list = None,  # سازگاری با app.py
+        max_new_tokens: int = 200
+    ) -> Tuple[str, Dict[str, str]]:
+        """
+        ناشناس‌سازی متن
+        
+        Args:
+            text: متن ورودی
+            entity_types: لیست انواع موجودیت‌ها
+            entities_to_anonymize: لیست انواع موجودیت‌ها (سازگاری)
+            max_new_tokens: حداکثر توکن خروجی
+            
+        Returns:
+            (متن ناشناس شده, mapping dictionary)
+        """
+        if not self._model_loaded:
+            logger.error("❌ مدل لود نشده است!")
+            return text, {}
+        
+        if not text or not text.strip():
+            return "", {}
+        
+        # سازگاری با app.py
+        if entities_to_anonymize is not None:
+            entity_types = entities_to_anonymize
+        
+        # تنظیم entity types پیش‌فرض
+        if entity_types is None:
+            entity_types = ["person", "company", "amount", "percent"]
+        
+        # ساخت prompt
+        prompt = self._create_prompt(text, entity_types)
+        
+        # تولید خروجی
+        try:
+            generated = self._generate(prompt, max_new_tokens)
+            
+            # پردازش خروجی
+            anonymized_text, mapping = self._parse_output(generated)
+            
+            return anonymized_text, mapping
+            
+        except Exception as e:
+            logger.error(f"❌ خطا در anonymization: {e}")
+            return text, {}
+    
+    def _create_prompt(self, text: str, entity_types: list) -> str:
+        """ساخت prompt برای مدل"""
+        # تبدیل entity types به فرمت مناسب
+        entity_mappings = {
+            "person": "اسامی اشخاص → person-01, person-02, ...",
+            "company": "نام شرکت‌ها → company-01, company-02, ...",
+            "amount": "اعداد و مبالغ → amount-01, amount-02, ...",
+            "percent": "درصدها → percent-01, percent-02, ..."
+        }
+        
+        instructions = [entity_mappings.get(et, et) for et in entity_types]
+        instructions_text = "\n".join([f"{i+1}. {inst}" for i, inst in enumerate(instructions)])
+        
+        prompt = f"""<|im_start|>system
+شما یک سیستم هوش مصنوعی برای ناشناس‌سازی متون فارسی هستید.
+<|im_end|>
+<|im_start|>user
+متن زیر را ناشناس کنید:
+{instructions_text}
+
+متن:
+{text}
+
+خروجی: فقط متن ناشناس شده
+<|im_end|>
+<|im_start|>assistant
+"""
+        return prompt
+    
+    def _generate(self, prompt: str, max_tokens: int = 200) -> str:
+        """تولید متن با مدل"""
+        try:
+            # Tokenize
+            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+            
+            # Generate با پارامترهای بهینه شده
+            logger.info("🔄 در حال تولید متن...")
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    **inputs,
+                    max_new_tokens=max_tokens,
+                    min_new_tokens=20,  # ⭐ حداقل طول خروجی
+                    do_sample=False,  # greedy decoding
+                    repetition_penalty=1.2,  # کاهش penalty
+                    pad_token_id=self.tokenizer.pad_token_id,
+                    eos_token_id=self.tokenizer.eos_token_id,
+                    num_beams=1,  # greedy search
+                    temperature=1.0,  # ⭐ اضافه شد
+                    top_p=None,  # ⭐ غیرفعال کردن sampling
+                    top_k=None   # ⭐ غیرفعال کردن sampling
+                )
+            
+            # Decode
+            result = self.tokenizer.decode(
+                outputs[0][inputs['input_ids'].shape[1]:],
+                skip_special_tokens=True
+            ).strip()
+            
+            logger.info(f"✅ تولید متن موفق - طول: {len(result)} کاراکتر")
+            logger.info(f"📝 خروجی خام: {result[:200]}...")  # ⭐ نمایش خروجی
+            return result
+                
+        except Exception as e:
+            logger.error(f"❌ خطا در تولید: {e}")
+            raise
+    
+    def _parse_output(self, output: str) -> Tuple[str, Dict[str, str]]:
+        """پردازش خروجی مدل"""
+        try:
+            logger.info(f"🔍 پردازش خروجی - طول: {len(output)} کاراکتر")
+            logger.info(f"📝 خروجی اولیه: {output[:300]}...")
+            
+            # تمیز کردن خروجی از prompt های تکراری
+            if "متن زیر را ناشناس کنید" in output:
+                output = output.split("متن زیر را ناشناس کنید")[0].strip()
+                logger.info("🧹 پاک کردن prompt تکراری")
+            
+            # حذف "خروجی:" و متن بعدش
+            if "خروجی:" in output:
+                # بعضی وقتا مدل می‌نویسه "خروجی: فقط متن ناشناس شده"
+                # و بعد متن اصلی رو میده
+                parts = output.split("خروجی:")
+                if len(parts) > 1:
+                    # اگه بعد از "خروجی:" متن داریم، اونو بگیر
+                    output = parts[1].strip()
+                    # پاک کردن "فقط متن ناشناس شده" اگه ��ست
+                    output = output.replace("فقط متن ناشناس شده", "").strip()
+                    logger.info("🧹 پاک کردن 'خروجی:'")
+            
+            # حذف newline های اضافی - ولی نه همه!
+            lines = output.split("\n")
+            if len(lines) > 3:  # اگه خیلی زیاد newline داره
+                # فقط اولین چند خط رو نگه دار
+                output = "\n".join(lines[:3]).strip()
+            
+            # حذف خطوط که شامل "(iParam" یا چیزهای غیرمرتبط هستن
+            lines = output.split("\n")
+            clean_lines = []
+            for line in lines:
+                line = line.strip()
+                # خطوط مفید رو نگه دار
+                if line and not any(x in line for x in ["(iParam", "متن زیر", "خروجی:", "###"]):
+                    clean_lines.append(line)
+            
+            if clean_lines:
+                anonymized_text = " ".join(clean_lines[:2])  # حداکثر 2 خط اول
+            else:
+                anonymized_text = output.strip()
+            
+            logger.info(f"✅ متن تمیز شده: {anonymized_text[:200]}...")
+            
+            # استخراج mapping از توکن‌ها
+            mapping = {}
+            # پیدا کردن همه توکن‌ها
+            tokens = re.findall(r'(person|company|amount|percent)-(\d+)', anonymized_text)
+            
+            logger.info(f"🔍 توکن‌های پیدا شده: {tokens}")
+            
+            for entity_type, number in tokens:
+                token = f"{entity_type}-{number}"
+                if token not in mapping:
+                    mapping[token] = f"[{token}]"  # placeholder
+            
+            logger.info(f"✅ پردازش موفق: {len(mapping)} توکن استخراج شد")
+            logger.info(f"📋 Mapping: {mapping}")
+            
+            return anonymized_text, mapping
+            
+        except Exception as e:
+            logger.error(f"❌ خطا در parse: {e}")
+            return output, {}
+    
+    def deanonymize(self, anonymized_text: str, mapping: Dict[str, str]) -> str:
+        """بازگردانی متن ناشناس شده"""
+        result = anonymized_text
+        
+        # معکوس کردن mapping
+        reverse_mapping = {v: k for k, v in mapping.items()}
+        
+        # جایگزینی توکن‌ها با مقادیر اصلی
+        for token, original in reverse_mapping.items():
+            result = result.replace(token, original)
+        
+        return result
+
+
+def create_qwen_anonymizer(
+    model_path: str = "./qwen-anonymizer-v2",
+    device: str = "auto"
+) -> QwenAnonymizer:
+    """
+    ایجاد instance از QwenAnonymizer
+    
+    Args:
+        model_path: مسیر مدل در Space
+        device: دستگاه (auto, cuda, cpu)
+    
+    Returns:
+        QwenAnonymizer instance
+    """
+    return QwenAnonymizer(model_path=model_path, device=device)
+
+
+if __name__ == "__main__":
+    # تست
+    print("=" * 60)
+    print("🧪 تست QwenAnonymizer")
+    print("=" * 60)
+    
+    try:
+        anonymizer = create_qwen_anonymizer()
+        
+        test_text = "شرکت پتروشیمی با سرمایه 100 میلیارد ریال توسط علی احمدی تاسیس شد."
+        print(f"\n📝 متن ورودی:\n   {test_text}")
+        
+        result, mapping = anonymizer.anonymize(test_text)
+        print(f"\n🔒 متن ناشناس شده:\n   {result}")
+        print(f"\n📋 Mapping:\n   {json.dumps(mapping, ensure_ascii=False, indent=2)}")
+        
+    except Exception as e:
+        print(f"\n❌ خطا: {e}")