Spaces:
Paused
Paused
| """Model loading utilities for the Arabic teaching multi-agent system.""" | |
| import json | |
| import logging | |
| import os | |
| from pathlib import Path | |
| import torch | |
| from dotenv import load_dotenv | |
| from huggingface_hub import hf_hub_download | |
| from peft import PeftModel | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| load_dotenv() | |
| logger = logging.getLogger(__name__) | |
| FINETUNED_7B_TEACHING_PATH_LOCAL = "models/qwen-7b-arabic-teaching" | |
| FINETUNED_7B_GRADING_PATH_LOCAL = "models/qwen-7b-arabic-grading" | |
| FINETUNED_7B_TEACHING_PATH_HF = "kdiabagate/qwen-7b-arabic-teaching" | |
| FINETUNED_7B_GRADING_PATH_HF = "kdiabagate/qwen-7b-arabic-grading" | |
| BASE_7B_MODEL = "Qwen/Qwen2.5-7B-Instruct" | |
| def _load_adapter_config(model_path: str, use_hub: bool, token: str) -> dict: | |
| """Load adapter_config.json from Hub or local path.""" | |
| if use_hub: | |
| config_path = hf_hub_download( | |
| repo_id=model_path, filename="adapter_config.json", token=token | |
| ) | |
| else: | |
| config_path = Path(model_path) / "adapter_config.json" | |
| with open(config_path) as f: | |
| return json.load(f) | |
| def _load_finetuned_model( | |
| model_type: str, | |
| hf_path: str, | |
| local_path: str, | |
| use_hub: bool = True, | |
| ) -> tuple[AutoModelForCausalLM, AutoTokenizer]: | |
| """Load fine-tuned model (teaching or grading). | |
| Args: | |
| model_type: Model type for logging ("teaching" or "grading") | |
| hf_path: HuggingFace Hub path | |
| local_path: Local filesystem path | |
| use_hub: If True, load from Hub; else local | |
| Returns: | |
| Tuple of (model, tokenizer) | |
| Raises: | |
| FileNotFoundError: If local model not found | |
| RuntimeError: If model loading fails | |
| """ | |
| if use_hub: | |
| model_path = hf_path | |
| logger.info(f"Loading {model_type} model from HuggingFace Hub: {model_path}...") | |
| else: | |
| model_path_obj = Path(local_path) | |
| if not model_path_obj.exists(): | |
| raise FileNotFoundError( | |
| f"{model_type.title()} model not found at {local_path}. " | |
| "Use use_hub=True or train the model first." | |
| ) | |
| model_path = str(model_path_obj) | |
| logger.info(f"Loading {model_type} model from local path: {model_path}...") | |
| try: | |
| hf_token = os.getenv("HF_TOKEN") | |
| tokenizer = AutoTokenizer.from_pretrained( | |
| model_path, trust_remote_code=True, token=hf_token | |
| ) | |
| adapter_config = _load_adapter_config(model_path, use_hub, hf_token) | |
| base_model_name = adapter_config["base_model_name_or_path"] | |
| logger.info(f"Base model: {base_model_name}") | |
| logger.info("Loading base model...") | |
| base_model = AutoModelForCausalLM.from_pretrained( | |
| base_model_name, | |
| device_map="auto", | |
| trust_remote_code=True, | |
| token=hf_token, | |
| ) | |
| logger.info("Loading LoRA adapter...") | |
| model = PeftModel.from_pretrained(base_model, model_path, device_map="auto", token=hf_token) | |
| memory_gb = torch.cuda.memory_allocated() / 1e9 if torch.cuda.is_available() else 0 | |
| logger.info( | |
| f"✓ {model_type.title()} model loaded successfully (memory: ~{memory_gb:.1f}GB)" | |
| ) | |
| return model, tokenizer | |
| except Exception as e: | |
| raise RuntimeError(f"Failed to load {model_type} model: {e}") from e | |
| def load_teaching_model(use_hub: bool = True) -> tuple[AutoModelForCausalLM, AutoTokenizer]: | |
| """Load fine-tuned Qwen2.5-7B teaching model. | |
| Fine-tuned on 153 multi-turn conversations for warm teaching tone. | |
| """ | |
| return _load_finetuned_model( | |
| model_type="teaching", | |
| hf_path=FINETUNED_7B_TEACHING_PATH_HF, | |
| local_path=FINETUNED_7B_TEACHING_PATH_LOCAL, | |
| use_hub=use_hub, | |
| ) | |
| def load_grading_model(use_hub: bool = True) -> tuple[AutoModelForCausalLM, AutoTokenizer]: | |
| """Load fine-tuned Qwen2.5-7B grading model. | |
| Fine-tuned for flexible grading with synonym/typo handling. | |
| """ | |
| return _load_finetuned_model( | |
| model_type="grading", | |
| hf_path=FINETUNED_7B_GRADING_PATH_HF, | |
| local_path=FINETUNED_7B_GRADING_PATH_LOCAL, | |
| use_hub=use_hub, | |
| ) | |
| def load_all_models() -> dict: | |
| """Load all models for the multi-agent system.""" | |
| logger.info("Loading all models for multi-agent system...") | |
| try: | |
| teaching_model, teaching_tokenizer = load_teaching_model() | |
| grading_model, grading_tokenizer = load_grading_model() | |
| total_memory = torch.cuda.memory_allocated() / 1e9 if torch.cuda.is_available() else 0 | |
| logger.info(f"✓ All models loaded successfully (total memory: ~{total_memory:.1f}GB)") | |
| return { | |
| "teaching_model": teaching_model, | |
| "teaching_tokenizer": teaching_tokenizer, | |
| "grading_model": grading_model, | |
| "grading_tokenizer": grading_tokenizer, | |
| } | |
| except Exception as e: | |
| raise RuntimeError(f"Failed to load models: {e}") from e | |