File size: 2,616 Bytes
d25492d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
---
license: mit
---
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
import torch
import numpy as np
# Define the model and tokenizer
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=2)
# Define the key words and their corresponding labels
key_words = ['ascites', 'cirrhosis', 'liver disease']
labels = [0, 1]
# Define a function to preprocess the input text
def preprocess_text(text):
inputs = tokenizer.encode_plus(
text,
add_special_tokens=True,
max_length=512,
return_attention_mask=True,
return_tensors='pt'
)
return inputs
# Define a function to make predictions
def make_prediction(text):
inputs = preprocess_text(text)
outputs = model(inputs['input_ids'], attention_mask=inputs['attention_mask'])
logits = outputs.logits
probabilities = torch.nn.functional.softmax(logits, dim=1)
predicted_class = torch.argmax(probabilities)
return predicted_class.item()
# Define a function to get the clinic that the referral should be directed to
def get_clinic(text):
predicted_class = make_prediction(text)
if predicted_class == 1:
return 'Liver Clinic'
else:
return 'Kidney Clinic'
# Define the model's configuration
model_config = {
'model_type': 'distilbert',
'num_labels': 2,
'key_words': key_words,
'labels': labels
}
# Define the model's metadata
model_metadata = {
'name': 'Referral Clinic Classifier',
'description': 'A model that classifies referrals to either the Liver Clinic or Kidney Clinic based on the presence of certain key words.',
'author': 'Your Name',
'version': '1.0'
}
# Train the model
train_data = [
('Patient has ascites and cirrhosis.', 1),
('Patient has liver disease.', 1),
('Patient has kidney disease.', 0),
('Patient has liver failure.', 1),
('Patient has kidney failure.', 0),
]
for text, label in train_data:
inputs = preprocess_text(text)
labels = torch.tensor(label)
outputs = model(inputs['input_ids'], attention_mask=inputs['attention_mask'], labels=labels)
loss = outputs.loss
model.zero_grad()
loss.backward()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
optimizer.step()
# Save the model to a file
torch.save(model.state_dict(),'referral_clinic_classifier.pth')
with open('model_config.json', 'w') as f:
json.dump(model_config, f)
with open('model_metadata.json', 'w') as f:
json.dump(model_metadata, f) |