Upload .\src\models\bwenge_model.py with huggingface_hub

.//src//models//bwenge_model.py

@@ -0,0 +1,157 @@
+"""BwengeAi model architecture."""
+
+import logging
+from typing import Any
+
+import torch
+import torch.nn as nn
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    PreTrainedModel,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class BwengeModel:
+    """Manager for BwengeAi model."""
+
+    def __init__(self, config: dict[str, Any]):
+        self.config = config
+        self.model_config = config.get("model", {})
+        self.tokenizer = None
+        self.model = None
+
+    def load_base_model(
+        self,
+        model_name: str | None = None,
+        use_quantization: bool = False,
+        device_map: str | None = None,
+    ) -> tuple[PreTrainedModel, AutoTokenizer]:
+        """Load a base model for fine-tuning."""
+        if model_name is None:
+            model_name = self.model_config.get("base_model", "meta-llama/Llama-3.2-1B")
+        
+        logger.info(f"Loading base model: {model_name}")
+        
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        
+        use_cuda = torch.cuda.is_available()
+        
+        if device_map is None:
+            device_map = "auto" if use_cuda else None
+        
+        model_kwargs: dict[str, Any] = {
+            "trust_remote_code": True,
+        }
+        
+        if use_quantization:
+            bnb_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_compute_dtype=torch.float16,
+                bnb_4bit_use_double_quant=True,
+            )
+            model_kwargs["quantization_config"] = bnb_config
+            model_kwargs["device_map"] = "auto"
+        elif use_cuda:
+            model_kwargs["dtype"] = torch.float16
+            model_kwargs["device_map"] = "auto"
+        else:
+            model_kwargs["dtype"] = torch.float32
+            model_kwargs["device_map"] = None
+            model_kwargs["low_cpu_mem_usage"] = True
+        
+        model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
+        
+        self.tokenizer = tokenizer
+        self.model = model
+        
+        logger.info(f"Model loaded: {model.num_parameters():,} parameters")
+        return model, tokenizer
+
+    def setup_lora(self, model: PreTrainedModel) -> PreTrainedModel:
+        """Set up LoRA for parameter-efficient fine-tuning."""
+        from peft import LoraConfig, get_peft_model, TaskType
+        
+        lora_config = self.model_config.get("finetune", {}).get("lora", {})
+        
+        config = LoraConfig(
+            task_type=TaskType.CAUSAL_LM,
+            r=lora_config.get("r", 16),
+            lora_alpha=lora_config.get("lora_alpha", 32),
+            lora_dropout=lora_config.get("lora_dropout", 0.05),
+            target_modules=lora_config.get("target_modules", [
+                "q_proj", "k_proj", "v_proj", "o_proj",
+                "gate_proj", "up_proj", "down_proj",
+            ]),
+            bias="none",
+        )
+        
+        model = get_peft_model(model, config)
+        
+        trainable, total = model.get_nb_trainable_parameters()
+        logger.info(f"LoRA applied: {trainable:,} trainable / {total:,} total ({100 * trainable / total:.2f}%)")
+        
+        return model
+
+    def save_model(self, output_dir: str) -> None:
+        """Save the model and tokenizer."""
+        if self.model is None or self.tokenizer is None:
+            logger.error("No model loaded to save")
+            return
+        
+        logger.info(f"Saving model to {output_dir}")
+        self.model.save_pretrained(output_dir)
+        self.tokenizer.save_pretrained(output_dir)
+        logger.info("Model saved successfully")
+
+    def load_saved_model(self, model_dir: str) -> tuple[PreTrainedModel, AutoTokenizer]:
+        """Load a previously saved model."""
+        logger.info(f"Loading saved model from {model_dir}")
+        
+        tokenizer = AutoTokenizer.from_pretrained(model_dir)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_dir,
+            torch_dtype=torch.float16,
+            device_map="auto",
+        )
+        
+        self.tokenizer = tokenizer
+        self.model = model
+        
+        return model, tokenizer
+
+    def generate(
+        self,
+        prompt: str,
+        max_new_tokens: int = 256,
+        temperature: float = 0.7,
+        top_p: float = 0.9,
+        top_k: int = 50,
+    ) -> str:
+        """Generate text from a prompt."""
+        if self.model is None or self.tokenizer is None:
+            logger.error("No model loaded for generation")
+            return ""
+        
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+        
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+                do_sample=True,
+                pad_token_id=self.tokenizer.eos_token_id,
+            )
+        
+        generated = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        return generated[len(prompt):]