FPRT/SurgerySort

1import streamlit as st2import numpy as np3import evaluate4from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification, TrainingArguments, Trainer, DataCollatorWithPadding5 6# Load tokenizer and model7tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")8model = AutoModel.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")9 10# Define label mappings11id2label = {0: "SURGERY", 1: "NON-SURGERY"}12label2id = {"SURGERY": 0, "NON-SURGERY": 1}13 14# Load evaluation metric15accuracy = evaluate.load("accuracy")16 17# Define preprocessing function18def preprocess_function(examples):19    return tokenizer(examples, truncation=True, padding=True)20 21# Load model for sequence classification22model = AutoModelForSequenceClassification.from_pretrained(23    "emilyalsentzer/Bio_ClinicalBERT", num_labels=2, id2label=id2label, label2id=label2id24)25 26# Define compute_metrics function27def compute_metrics(eval_pred):28    predictions, labels = eval_pred29    predictions = np.argmax(predictions, axis=1)30    return accuracy.compute(predictions=predictions, references=labels)31 32# Define data collator33data_collator = DataCollatorWithPadding(tokenizer=tokenizer)34 35# Define training arguments36training_args = TrainingArguments(37    output_dir="my_awesome_model",38    learning_rate=2e-5,39    per_device_train_batch_size=16,40    per_device_eval_batch_size=16,41    num_train_epochs=2,42    weight_decay=0.01,43    evaluation_strategy="epoch",44    save_strategy="epoch",45    load_best_model_at_end=True,46    push_to_hub=True,47)48 49# Initialize trainer50trainer = Trainer(51    model=model,52    args=training_args,53    tokenizer=tokenizer,54    data_collator=data_collator,55    compute_metrics=compute_metrics,56)57 58# Streamlit UI59st.title("Clinical Text Classification")60text = st.text_area("Enter clinical text:", "")61 62if st.button("Classify"):63    # Tokenize user input and predict64    tokenized_text = preprocess_function(text)65    result = trainer.predict(tokenized_text)66    prediction = np.argmax(result.predictions, axis=1)[0]67    st.write("Predicted Label:", id2label[prediction])68