Fixed the Python implementation in the README

README.md

@@ -27,18 +27,47 @@ This model is designed for integration into IT call center software systems, whe
 2. Load the model using the following code snippet:
 
 ```python
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from __future__ import annotations
+from transformers import RobertaConfig, RobertaModel, RobertaTokenizer, AutoModel, AutoTokenizer
 
-model_name = "KameronB/sitcc-roberta"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
+# Add a custom regression head to RoBERTa
+class SITCC(torch.nn.Module):
+    def __init__(self, model):
+        super(SITCC, self).__init__()
+        self.roberta = model
+        self.regressor = torch.nn.Linear(config.hidden_size, 1)  # Outputs a single value
+        
 
-def predict(sentence):
-    inputs = tokenizer(sentence, return_tensors="pt")
-    outputs = model(**inputs)
-    return outputs.logits
+    def forward(self, input_ids, attention_mask):
+        outputs = self.roberta(input_ids=input_ids, attention_mask=attention_mask)
+        sequence_output = outputs[1]  # The last hidden-state is the first element of the output tuple
+        logits = self.regressor(sequence_output)
+        return logits
 
-sentence = "Your example sentence here."
-score = predict(sentence)
-print(f"Quality Score: {score}")
+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)
+
+    # create the model based on the RoBERTa base model
+    model = SITCC(RobertaModel(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
+
+
+model = 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