Views
No views yet
python 3.9 or higher, torch 2.0.0 or higher, transformers 4.31.0 or higher.git clone https://github.com/anti-aii/RagE.git
cd RagE
pip install -e .1>>> import torch
2>>> from pyvi import ViTokenizer
3>>> from rage import SentenceEmbedding
4>>> device= torch.device('cuda' if torch.cuda.is_available() else 'cpu')
5>>> model= SentenceEmbedding(model_name= "vinai/phobert-base-v2", torch_dtype= torch.float32, aggregation_hidden_states= False, strategy_pooling= "dense_first")
6>>> model.to(device)
7SentenceEmbeddingConfig(model_base: {'model_type_base': 'RobertaModel', 'model_name': 'vinai/phobert-base-v2', 'type_backbone': 'mlm', 'required_grad_base_model': True, 'aggregation_hidden_states': False, 'concat_embeddings': False, 'dropout': 0.1, 'quantization_config': None}, pooling: {'strategy_pooling': 'dense_first'})1>>> model.summary_params()
2trainable params: 135588864 || all params: 135588864 || trainable%: 100.0
3>>> model.summary()
4+---------------------------+-------------+------------------+
5| Layer (type) | Params | Trainable params |
6+---------------------------+-------------+------------------+
7| model (RobertaModel) | 134,998,272 | 134998272 |
8| pooling (PoolingStrategy) | 590,592 | 590592 |
9| drp1 (Dropout) | 0 | 0 |
10+---------------------------+-------------+------------------+1>>> model.load("best_sup_general_embedding_phobert2.pt", key= False)
2>>> sentences= ["Tôi đang đi học", "Bạn tên là gì?",]
3>>> sentences= list(map(lambda x: ViTokenizer.tokenize(x), sentences))
4>>> model.encode(sentences, batch_size= 1, normalize_embedding= "l2", return_tensors= "np", verbose= 1)
52/2 [==============================] - 0s 43ms/Sample
6array([[ 0.00281098, -0.00829096, -0.01582766, ..., 0.00878178,
7 0.01830498, -0.00459659],
8 [ 0.00249859, -0.03076724, 0.00033016, ..., 0.01299141,
9 -0.00984358, -0.00703243]], dtype=float32)>>> model= SentenceEmbedding.from_pretrained('anti-ai/VieSemantic-base')1>>> sentences = ["Nó rất thú_vị", "Nó không thú_vị ."]
2>>> output= model.encode(sentences, batch_size= 1, return_tensors= 'pt')
3>>> torch.cosine_similarity(output[0].view(1, -1), output[1].view(1, -1)).cpu().tolist()
42/2 [==============================] - 0s 40ms/Sample
5[0.5605039596557617]| Model Name | Model Type | #params | checkpoint |
|---|---|---|---|
| anti-ai/ViEmbedding-base | SentenceEmbedding | 135.5M | model |
| anti-ai/BioViEmbedding-base-unsup | SentenceEmbedding | 135.5M | model |
| anti-ai/VieSemantic-base | SentenceEmbedding | 135.5M | model |