Spaces:

mickey1976
/

cove-api

Runtime error

cove-api / src /utils /embedding_injector.py

mickey1976

Deploy: Minimal FastAPI backend for CoVE Space

549c270 about 1 year ago

850 Bytes

# src/utils/embedding_injector.py

import torch.nn as nn
from transformers import GPT2LMHeadModel
from utils.emb_comp_utils import get_peft_embedding

def inject_hashed_embedding(base_model: GPT2LMHeadModel, tokenizer) -> GPT2LMHeadModel:
    """
    Replaces the base model's input embedding layer with a hashed embedding layer,
    using the get_peft_embedding function. Returns the modified model.
    """
    print("[ℹ] Injecting hashed embedding layer...")

    # Create hashed embedding layer
    new_embedding = get_peft_embedding(
        vocab_size=tokenizer.vocab_size,
        hidden_size=base_model.config.hidden_size,
        num_hashes=2,
        num_buckets=8192,  # or 4096 if needed
        device=base_model.device
    )

    # Replace original embedding layer
    base_model.transformer.wte = new_embedding

    return base_model

	# src/utils/embedding_injector.py

	import torch.nn as nn
	from transformers import GPT2LMHeadModel
	from utils.emb_comp_utils import get_peft_embedding

	def inject_hashed_embedding(base_model: GPT2LMHeadModel, tokenizer) -> GPT2LMHeadModel:
	"""
	Replaces the base model's input embedding layer with a hashed embedding layer,
	using the get_peft_embedding function. Returns the modified model.
	"""
	print("[ℹ] Injecting hashed embedding layer...")

	# Create hashed embedding layer
	new_embedding = get_peft_embedding(
	vocab_size=tokenizer.vocab_size,
	hidden_size=base_model.config.hidden_size,
	num_hashes=2,
	num_buckets=8192, # or 4096 if needed
	device=base_model.device
	)

	# Replace original embedding layer
	base_model.transformer.wte = new_embedding

	return base_model