Update README.md

README.md

@@ -9,7 +9,8 @@ To do:
 4. Loss function.
 
 To run TE_Embedding model:
-`import os
+```python
+import os
 from transformers import (AutoConfig,
     AutoTokenizer,AutoModelForCausalLM
 )
@@ -30,7 +31,7 @@ class TEmbeddingModel(torch.nn.Module):
             [torch.nn.Linear(self.hidden_size, self.hidden_size//len(self.prompt_suffixes))
              for _ in range(len(self.prompt_suffixes))])
         self.tokenizer, self.llama = self.load_llama()
-        self.device = torch.device('cuda')
+        # self.device = torch.device('cuda')
         self.tanh = torch.nn.Tanh()
         self.suffixes_ids = []
         self.suffixes_ids_len = []
@@ -79,12 +80,12 @@ class TEmbeddingModel(torch.nn.Module):
         suffixes_ones = self.suffixes_ones.unsqueeze(0) 
         suffixes_ones = suffixes_ones.repeat(batch_size, 1)
         device = next(self.parameters()).device
-        attention_mask = torch.cat([attention_mask, suffixes_ones], dim=-1).to('cuda')
+        attention_mask = torch.cat([attention_mask, suffixes_ones], dim=-1).to(device)
 
         suffixes_ids = self.suffixes_ids.unsqueeze(0)
         suffixes_ids = suffixes_ids.repeat(batch_size, 1)
-        input_ids = torch.cat([input_ids, suffixes_ids], dim=-1).to('cuda')
-        last_hidden_state = self.llama.base_model(attention_mask=attention_mask, input_ids=input_ids).last_hidden_state.to('cuda')
+        input_ids = torch.cat([input_ids, suffixes_ids], dim=-1) #to("cuda")
+        last_hidden_state = self.llama.base_model(attention_mask=attention_mask, input_ids=input_ids).last_hidden_state.to(device)
         index = -1
         for i in range(len(self.suffixes_ids_len)):
             embedding = last_hidden_state[:, index, :]
@@ -119,4 +120,5 @@ if __name__ == "__main__":
         output = TE_model(["Hello", "Nice to meet you"])
         cos_sim = F.cosine_similarity(output[0],output[1],dim=0)
         print(cos_sim)
-`
+
+```