Update README.md

README.md

@@ -29,14 +29,14 @@ This model is designed for integration into IT call center software systems, whe
 ```python
 from __future__ import annotations
 from transformers import RobertaConfig, RobertaModel, RobertaTokenizer, AutoModel, AutoTokenizer
+import torch
 
 # Add a custom regression head to RoBERTa
 class SITCC(torch.nn.Module):
-    def __init__(self, model):
+    def __init__(self, model, config):
         super(SITCC, self).__init__()
         self.roberta = model
         self.regressor = torch.nn.Linear(config.hidden_size, 1)  # Outputs a single value
-        
 
     def forward(self, input_ids, attention_mask):
         outputs = self.roberta(input_ids=input_ids, attention_mask=attention_mask)
@@ -44,30 +44,35 @@ class SITCC(torch.nn.Module):
         logits = self.regressor(sequence_output)
         return logits
 
+
 def init_model() -> SITCC:
     # Load the model from huggingface
     model_name = "KameronB/sitcc-roberta"
     tokenizer = AutoTokenizer.from_pretrained(model_name, from_tf=False)
-    config = RobertaConfig.from_pretrained(model_name)
-
+    config = RobertaConfig.from_pretrained(model_name,)
+ 
     # create the model based on the RoBERTa base model
-    model = SITCC(RobertaModel(config))
+    model = SITCC(RobertaModel(config), config)
+
     # fetch the statedict to apply the fine-tuned weights
     state_dict = torch.hub.load_state_dict_from_url(f"https://huggingface.co/{model_name}/resolve/main/pytorch_model.bin")
-    model.load_state_dict(state_dict)
-    return model
-
+    # if running on cpu
+    # state_dict = torch.hub.load_state_dict_from_url(f"https://huggingface.co/{model_name}/resolve/main/pytorch_model.bin", map_location=torch.device('cpu'))
 
-model = init_model()
+    model.load_state_dict(state_dict)
+    return model, tokenizer
 
+model, tokenizer = init_model()
 
 def predict(sentences):
+
     model.eval()
     inputs = tokenizer(sentences, padding=True, truncation=True, max_length=512, return_tensors="pt")
     input_ids = inputs['input_ids']
     attention_mask = inputs['attention_mask']
-    
+
     with torch.no_grad():
         outputs = model(input_ids, attention_mask)
+
     return outputs