File size: 2,667 Bytes
7d9df81 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
import os
import torch
import torch.nn as nn
import numpy as np
import random
from transformers import (
BartForConditionalGeneration,
AutoModelForCausalLM,
BertModel,
Wav2Vec2Model,
CLIPModel,
AutoTokenizer
)
class MultiModalModel(nn.Module):
def __init__(self):
super(MultiModalModel, self).__init__()
# 初始化子模型
self.text_generator = BartForConditionalGeneration.from_pretrained('facebook/bart-base')
self.code_generator = AutoModelForCausalLM.from_pretrained('gpt2')
self.nlp_encoder = BertModel.from_pretrained('bert-base-uncased')
self.speech_encoder = Wav2Vec2Model.from_pretrained('facebook/wav2vec2-base-960h')
self.vision_encoder = CLIPModel.from_pretrained('openai/clip-vit-base-patch32')
# 初始化分词器和处理器
self.text_tokenizer = AutoTokenizer.from_pretrained('facebook/bart-base')
self.code_tokenizer = AutoTokenizer.from_pretrained('gpt2')
self.nlp_tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
self.speech_processor = AutoTokenizer.from_pretrained('facebook/wav2vec2-base-960h')
self.vision_processor = AutoTokenizer.from_pretrained('openai/clip-vit-base-patch32')
def forward(self, task, inputs):
if task == 'text_generation':
# 确保 attention_mask 在 inputs 中
attention_mask = inputs.get('attention_mask')
print("输入数据:", inputs)
outputs = self.text_generator.generate(
inputs['input_ids'],
max_new_tokens=100, # 增加生成的最大新令牌数
pad_token_id=self.text_tokenizer.eos_token_id,
attention_mask=attention_mask,
top_p=0.9, # 调整 top_p 值
top_k=50, # 保持 top_k 值
temperature=0.8, # 调整 temperature 值
do_sample=True
)
print("生成的输出:", outputs)
return self.text_tokenizer.decode(outputs[0], skip_special_tokens=True)
# 根据需要添加其他任务的逻辑...
# 主函数
if __name__ == "__main__":
# 初始化模型
model = MultiModalModel()
# 示例任务和输入数据
task = "text_generation"
input_text = "This is a sample input."
tokenizer = model.text_tokenizer
inputs = tokenizer(input_text, return_tensors='pt')
# 添加 attention_mask 键值对
inputs['attention_mask'] = torch.ones_like(inputs['input_ids'])
# 模型推理
result = model(task, inputs)
print("最终输出结果：", result)
trust_remote_code=True
|