Merge pull request #6 from Yiran-ASU/yiran-fix

Fixed the segmentation fault in CamemBERT and VIT

experiments/CamemBERT/model/CamemBERT.py

@@ -24,8 +24,9 @@ def __init__(self, model_name: str):
         self.layer3 = AutoModelForTokenClassification.from_pretrained(model_name).classifier
 
     def forward(self, input):
+        x = input.reshape(1, 7, 768)
         # Return only the logits.
-        x = self.layer1(input).last_hidden_state
+        x = self.layer1(x).last_hidden_state
         x = self.layer2(x)
         x = self.layer3(x)
         return x
@@ -44,7 +45,11 @@ def forward(self, input):
 
 print("Parsing sentence tokens.")
 example_input = prepare_sentence_tokens(model_name, sentence)
-print(example_input.shape)
+print("example_input shape: ", example_input.shape)
+
+# The original example_input shape is [1, 7, 768], now we reshape it into [1, 7*768]
+example_input = example_input.reshape(1, 7*768)
+print("example_input shape after reshaping: ", example_input.shape)
 
 print("Instantiating model.")
 model = OnlyLogitsHuggingFaceModel(model_name)

experiments/VIT/model/vit.py

@@ -25,6 +25,7 @@ def __init__(self):
 		self.layer2 = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224').vit.layernorm
 		self.layer3 = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224').classifier
 	def forward(self,x):
+		x = x.reshape(1, 197, 768)
 		x = self.layer1(x).last_hidden_state
 		x = self.layer2(x)
 		x = self.layer3(x)
@@ -38,6 +39,11 @@ def forward(self,x):
 example_input = model(inputs)
 print(example_input.shape)
 
+# The original example_input shape is [1, 197, 768], now we reshape it into [1, 197*768]
+example_input = example_input.reshape(1, 197*768)
+print("example_input shape after reshaping: ", example_input.shape)
+
+
 vit_model = vit().eval()
 output = vit_model(example_input)
 print(output.shape)