Update README.md

README.md

@@ -20,12 +20,16 @@ Labels:
 ## Usage
 
 ``` python
-from transformers import AutoTokenizer
+from transformers import AutoModel, AutoTokenizer
+import torch
+from torch.nn import Linear, Module, TransformerEncoderLayer, CrossEntropyLoss
+from huggingface_hub import PyTorchModelHubMixin
+from typing import Optional
 
 HF_MODEL_NAME = 'poc-embeddings/rubert-tiny-turbo-godeal'
 MODEL_NAME = 'sergeyzh/rubert-tiny-turbo'
 
-id2label = {0:'noise', 1:'demand', 2:'noise'}
+id2label = {0:'noise', 1:'supply', 2:'demand'}
 
 class SupplyDemandTrader(
  Module,
@@ -92,9 +96,11 @@ model = SupplyDemandTrader.from_pretrained(HF_MODEL_NAME)
 tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_NAME)
 model.eval()
 
+message = "Куплю Iphone 8"
+
 with torch.inference_mode():
- ids = tokenizer("Куплю Iphone 8", return_tensors="pt")
- logits = checkpoint.forward(ids['input_ids'], ids['attention_mask']))
+ ids = tokenizer(message, return_tensors="pt")
+ logits = model.forward(ids['input_ids'], ids['attention_mask'])['logits']
  preds = torch.argmax(logits)
  print(id2label[int(preds)])
 ```