from models.Gemini import Gemini
from models.OpenAI import ChatGPT
from models.OpenAI import GPT4
from models.DeepSeek import DeepSeekModel
from models.Pangu import Pangu
from models.Qwen import Qwen
from models.Pangu72B import Pangu72B

class ModelFactory:
    @staticmethod
    def get_model_class(model_name):
        if model_name == "Gemini":
            return Gemini
        elif model_name == "ChatGPT":
            return ChatGPT
        elif model_name == "GPT4":
            return GPT4
        elif model_name == "DeepSeek":
            return DeepSeekModel
        elif model_name == "Pangu":
            return Pangu
        elif model_name == "Qwen":
            return Qwen
        elif model_name == "Pangu72B":
            return Pangu72B
        else:
            raise Exception(f"Unknown model name {model_name}")