from hyperformer.adapters import AdapterController, AutoAdapterConfig
from hyperformer.third_party.models import T5Config, T5ForConditionalGeneration
from transformers import AutoTokenizer, set_seed
import os
os.environ["CUDA_VISIBLE_DEVICES"]="0"
set_seed(42)
config = T5Config.from_pretrained('t5-3b',cache_dir="/local/nlpswordfish/tuhin/")
tokenizer = AutoTokenizer.from_pretrained('t5-3b',cache_dir="/local/nlpswordfish/tuhin/")
adapter_config = AutoAdapterConfig.get('meta-adapter')
#####################
adapter_config.input_dim = 1024
adapter_config.tasks = data_args.tasks
adapter_config.device = training_args.device
adapter_config.task_to_adapter = {task:adapter for task, adapter in zip(data_args.tasks, data_args.adapters)} if data_args.adapters is not None else None
adapter_config.task_to_embeddings = {task:embedding for task, embedding in zip(data_args.tasks, data_args.task_embeddings)} if (data_args.task_embeddings is not None) else None
######################
extra_adapter_params = ("task_embedding_dim","add_layer_norm_before_adapter","add_layer_norm_after_adapter","reduction_factor","hidden_dim","non_linearity","train_task_embeddings","projected_task_embedding_dim","task_hidden_dim","conditional_layer_norm","train_adapters_blocks","unique_hyper_net","unique_hyper_net_layer_norm","efficient_unique_hyper_net")
for p in extra_adapter_params:
if hasattr(adapter_args, p) and hasattr(adapter_config, p):
setattr(adapter_config, p, getattr(adapter_args, p))
model = T5ForConditionalGeneration.from_pretrained("/mnt/swordfish-datastore/tuhin/hyperformer++",from_tf=False, config=config,cache_dir="/local/nlpswordfish/tuhin/",adapter_config=adapter_config)
model.cuda()
inputs = tokenizer.encode("it 's a charming and often affecting journey .", return_tensors="pt")
gen_kwargs = {"max_length": 256, "num_beams": 1}
gen_kwargs["task"] = "sst"
gen_kwargs["task_embedding"] = model.task_embedding_controller("sst") if (self.config.train_adapters and isinstance(self.adapter_config, MetaAdapterConfig)) else None
outputs = model.generate(input_ids=inputs.cuda(),**gen_kwargs)
answer = tokenizer.decode(outputs[0],skip_special_tokens=True)
print("Predicted output", answer)