from transformers import pipeline
import os
import openai
import gc  # برای آزاد کردن حافظه
import psutil  # برای مانیتور کردن حافظه

class MultiModelSystem:
    def __init__(self):
        self.models = {}

    def load_model(self, task, model_id, use_auth_token=None):
        """مدل‌ها را به صورت lazy بارگذاری می‌کند."""
        if model_id == "OpenAI":
            return self.load_openai_model()
        else:
            return pipeline(
                task=self.get_task_type(task),
                model=model_id,
                use_auth_token=use_auth_token
            )

    def unload_model(self, task):
        """مدل بارگذاری شده را از حافظه پاک می‌کند."""
        if task in self.models:
            del self.models[task]
            gc.collect()  # جمع‌آوری حافظه

    @staticmethod
    def load_openai_model():
        """مدل ریاضی OpenAI."""
        return "OpenAI (Math)"

    @staticmethod
    def get_task_type(task):
        task_map = {
            "translation": "translation",
            "qa": "question-answering",
            "persian_nlp": "text-classification",
            "custom_ai": "text-generation",
            "math": "text-generation",
        }
        return task_map.get(task, "text-generation")

    def process_task(self, task, model_id, **kwargs):
        """مدیریت وظایف."""
        if task not in self.models:
            self.models[task] = self.load_model(task, model_id, use_auth_token=os.getenv("Passsssssss"))
        model = self.models[task]

        # پردازش وظیفه
        if task == "math":
            return self.process_math_task(kwargs.get("text"))
        else:
            return model(**kwargs)

    def process_math_task(self, text):
        """مدیریت وظایف ریاضی OpenAI."""
        try:
            response = openai.ChatCompletion.create(
                model="gpt-4",
                messages=[{"role": "user", "content": text}]
            )
            return response['choices'][0]['message']['content'].strip()
        except Exception as e:
            print(f"Error processing math task: {e}")
            return None

# مانیتورینگ حافظه
def check_memory_usage():
    mem = psutil.virtual_memory()
    print(f"Memory usage: {mem.percent}% ({mem.used / (1024 ** 3):.2f} GB used)")

# نمونه استفاده
if __name__ == "__main__":
    MODEL_CONFIG = {
        "translation": "PontifexMaximus/opus-mt-iir-en-finetuned-fa-to-en",
        "qa": "HooshvareLab/bert-fa-base-uncased",
        "math": "OpenAI",
        "persian_nlp": "HooshvareLab/bert-fa-zwnj-base",
        "custom_ai": "universitytehran/PersianMind-v1.0",
    }

    system = MultiModelSystem()

    tasks = [
        {"task": "translation", "model_id": MODEL_CONFIG["translation"], "kwargs": {"text": "سلام دنیا!", "src_lang": "fa", "tgt_lang": "en"}},
        {"task": "qa", "model_id": MODEL_CONFIG["qa"], "kwargs": {"question": "پایتخت ایران چیست؟", "context": "ایران کشوری در خاورمیانه است و پایتخت آن تهران است."}},
        {"task": "math", "model_id": MODEL_CONFIG["math"], "kwargs": {"text": "What is the integral of x^2?"}},
    ]

    for task_info in tasks:
        check_memory_usage()  # نمایش میزان حافظه
        result = system.process_task(task_info["task"], task_info["model_id"], **task_info["kwargs"])
        print(f"Result for task '{task_info['task']}':", result)
        system.unload_model(task_info["task"])  # آزادسازی حافظه
        check_memory_usage()  # دوباره بررسی حافظه