Create model_utils.py

model_utils.py

@@ -0,0 +1,39 @@
+# model_utils.py
+
+from transformers import (
+    AutoTokenizer,
+    AutoModel,
+    AutoModelForCausalLM,
+)
+import torch
+
+
+MODEL_OPTIONS = {
+    "BERT (bert-base-uncased)": "bert-base-uncased",
+    "DistilBERT": "distilbert-base-uncased",
+    "RoBERTa": "roberta-base",
+    "GPT-2": "gpt2",
+    "Electra": "google/electra-small-discriminator",
+    "ALBERT": "albert-base-v2",
+    "XLNet": "xlnet-base-cased",
+}
+
+
+def load_model(model_name):
+    if "gpt2" in model_name or "causal" in model_name:
+        model = AutoModelForCausalLM.from_pretrained(model_name, output_attentions=True)
+    else:
+        model = AutoModel.from_pretrained(model_name, output_attentions=True)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    return tokenizer, model
+
+
+def get_model_info(model):
+    config = model.config
+    return {
+        "Model Type": config.model_type,
+        "Number of Layers": getattr(config, "num_hidden_layers", "N/A"),
+        "Number of Attention Heads": getattr(config, "num_attention_heads", "N/A"),
+        "Total Parameters": sum(p.numel() for p in model.parameters()),
+    }