Update main.py

main.py

@@ -1,67 +1,69 @@
-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)
+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