Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForMaskedLM
3from sentence_transformers import SentenceTransformer
4from huggingface_hub import hf_hub_download
5
6# اطلاعات رپوزیتوری
7repo_id = "YasinProDebian/mental-health-fill-mask"
8
9# بارگذاری مدل و توکنایزر
10tokenizer = AutoTokenizer.from_pretrained(repo_id)
11model = AutoModelForMaskedLM.from_pretrained(repo_id)
12
13# بارگذاری Sentence-Transformer
14sentence_model = SentenceTransformer('paraphrase-multilingual-mpnet-base-v2')
15
16# تنظیم دستگاه (GPU یا CPU)
17device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
18model.to(device)
19sentence_model.to(device)
20
21# بارگذاری لایه projection
22projection_path = hf_hub_download(repo_id=repo_id, filename="projection_layer.pth")
23combined_dim = model.config.hidden_size + sentence_model.get_sentence_embedding_dimension() # 1024 + 768 = 1792
24hidden_size = model.config.hidden_size # 1024
25projection_layer = torch.nn.Linear(combined_dim, hidden_size)
26projection_layer.load_state_dict(torch.load(projection_path, map_location=device))
27projection_layer.to(device)
28
29def predict_top_k(text, k=10):
30 model.eval()
31 with torch.no_grad():
32 # توکنایز کردن متن
33 inputs = tokenizer(text, padding=True, truncation=True, max_length=128, return_tensors="pt").to(device)
34 sentence_embedding = sentence_model.encode([text], convert_to_tensor=True).to(device)
35
36 # پیدا کردن موقعیت <mask>
37 mask_token_index = torch.where(inputs["input_ids"][0] == tokenizer.mask_token_id)[0]
38 if len(mask_token_index) == 0:
39 return "متن باید شامل <mask> باشد."
40
41 mask_token_index = mask_token_index[0].item() # موقعیت اولین <mask>
42
43 # گرفتن خروجی از مدل RoBERTa
44 outputs = model.roberta(**inputs, output_hidden_states=True)
45 hidden_states = outputs.last_hidden_state
46
47 # ترکیب با امبدینگ Sentence-Transformer
48 sentence_embedding_expanded = sentence_embedding.unsqueeze(1).expand(-1, hidden_states.size(1), -1)
49 combined_input = torch.cat((hidden_states, sentence_embedding_expanded), dim=-1)
50
51 # اعمال لایه projection و پیشبینی نهایی
52 projected_output = projection_layer(combined_input)
53 logits = model.lm_head(projected_output)
54
55 # گرفتن احتمالات برای موقعیت <mask>
56 mask_logits = logits[0, mask_token_index] # فقط logits مربوط به <mask>
57 probs = torch.softmax(mask_logits, dim=-1) # تبدیل به احتمالات
58
59 # انتخاب k توکن برتر با احتمالاتشان
60 top_k_probs, top_k_indices = torch.topk(probs, k, dim=-1)
61
62 # تولید جملات و جمعآوری نتایج
63 results = []
64 input_ids = inputs["input_ids"][0].clone()
65 for prob, idx in zip(top_k_probs, top_k_indices):
66 input_ids[mask_token_index] = idx # جایگزینی <mask> با توکن پیشبینیشده
67 decoded_sentence = tokenizer.decode(input_ids, skip_special_tokens=True)
68 results.append({"sentence": decoded_sentence, "probability": prob.item()})
69
70 return results
71
72# تست تابع
73test_text = "ناتوانی در تمرکز و <mask> یک مشکل شایع است."
74top_predictions = predict_top_k(test_text, k=20)
75for i, pred in enumerate(top_predictions):
76 print(f"پیشبینی {i+1}: {pred['sentence']} (احتمال: {pred['probability']:.4f})")
77