Update README.md

README.md

@@ -31,19 +31,16 @@ First, ensure you have the required libraries installed:
 ```bash
 !pip install peft transformers sentence-transformers torch
 
-
-from transformers import AutoModel, AutoTokenizer
-from peft import PeftModel
-
-# Load the base model
-base_model_name = "sentence-transformers/all-mpnet-base-v2"
-base_model = AutoModel.from_pretrained(base_model_name)
-
-# Load the tokenizer
-tokenizer = AutoTokenizer.from_pretrained(base_model_name)
-
-# Load and apply LoRA weights
-lora_model = PeftModel.from_pretrained(base_model, "vahidthegreat/StanceAware-SBERT")
+from peft import LoraConfig, get_peft_model, prepare_model_for_int8_training, TaskType, PeftModel # peft-0.7.1
+from transformers import (
+    AutoModel,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    HfArgumentParser,
+    AutoTokenizer,
+    TrainingArguments,
+    AutoConfig,
+)
 ```
 
 
@@ -52,32 +49,39 @@ lora_model = PeftModel.from_pretrained(base_model, "vahidthegreat/StanceAware-SB
 
 The following custom `SiameseNetworkMPNet` class leverages the model for stance detection tasks. It pools embeddings and normalizes them for similarity calculations. This is for the sake of replicability of our exact results. But the model would work without this as well.
 ```bash
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-# Custom Siamese Network class
 class SiameseNetworkMPNet(nn.Module):
     def __init__(self, model_name, tokenizer, normalize=True):
         super(SiameseNetworkMPNet, self).__init__()
-        
-        # Initialize with LoRA-applied model
-        self.model = AutoModel.from_pretrained(model_name)
-        self.model = PeftModel.from_pretrained(self.model, "vahidthegreat/StanceAware-SBERT")
+
+        self.model = AutoModel.from_pretrained(model_name)#, quantization_config=bnb_config, trust_remote_code=True)
         self.normalize = normalize
         self.tokenizer = tokenizer
 
     def forward(self, **inputs):
         model_output = self.model(**inputs)
         attention_mask = inputs['attention_mask']
-        last_hidden_states = model_output.last_hidden_state  # Token embeddings
-        embeddings = torch.sum(last_hidden_states * attention_mask.unsqueeze(-1), 1) / torch.clamp(attention_mask.sum(1, keepdim=True), min=1e-9)  # Mean pooling
-
+        last_hidden_states = model_output.last_hidden_state  # First element of model_output contains all token embeddings
+        embeddings = torch.sum(last_hidden_states * attention_mask.unsqueeze(-1), 1) / torch.clamp(attention_mask.sum(1, keepdim=True), min=1e-9) # mean_pooling
         if self.normalize:
             embeddings = F.layer_norm(embeddings, embeddings.shape[1:])
             embeddings = F.normalize(embeddings, p=2, dim=1)
 
         return embeddings
+
+
+base_model_name = "sentence-transformers/all-mpnet-base-v2" 
+tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+
+# Load the base model
+base_model = SiameseNetworkMPNet(model_name=base_model_name, tokenizer=tokenizer)
+
+# Load and apply LoRA weights
+lora_model = SiameseNetworkMPNet(model_name=base_model_name, tokenizer=tokenizer)
+lora_model = PeftModel.from_pretrained(base_model_copy, "vahidthegreat/StanceAware-SBERT")
+lora_model = lora_model.merge_and_unload()
+
+base_model.eval()
+lora_model.eval()
 ```
 
 
@@ -105,9 +109,10 @@ def two_sentence_similarity(model, tokenizer, text1, text2):
 text1 = "I love pineapple on pizza"
 text2 = "I hate pineapple on pizza"
 
-# Instantiate model and tokenizer
-stance_model = SiameseNetworkMPNet(model_name=base_model_name, tokenizer=tokenizer)
-two_sentence_similarity(stance_model, tokenizer, text1, text2)
+print(f"For Base Model sentences: '{text1}' and '{text2}'")
+two_sentence_similarity(base_model, tokenizer, text1, text2)
+print(f"\n\nFor FineTuned Model sentences: '{text1}' and '{text2}'")
+two_sentence_similarity(lora_model, tokenizer, text1, text2)
 ```