Update app.py

app.py

@@ -1,20 +1,21 @@
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
-# Load a pre-trained version of ClinicalGPT
-model = AutoModelForCausalLM.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
-# Tokenize your clinical text data using the AutoTokenizer class
-tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
-# Convert your tokenized data into PyTorch tensors and create a PyTorch Dataset object
 import torch
-from torch.utils.data import Dataset
-
+from torch.utils.data import Dataset, random_split
+from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
 
+# Prompt user to enter clinical text data and corresponding labels
 train_texts = st.text_input("Enter your clinical text data (separated by commas):")
 train_labels = st.text_input("Enter your corresponding labels (separated by commas):")
 
+# Convert comma-separated values into lists
 train_texts = train_texts.split(",")
 train_labels = train_labels.split(",")
 
+# Load pre-trained model and tokenizer
+model = AutoModelForCausalLM.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
+tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
+
+# Create PyTorch Dataset object
 class ClinicalDataset(Dataset):
          def __init__(self, texts, labels, tokenizer):
              self.texts = texts
@@ -29,12 +30,15 @@ class ClinicalDataset(Dataset):
              label = self.labels[idx]
              encoding = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
              return {"input_ids": encoding["input_ids"].squeeze(), "attention_mask": encoding["attention_mask"].squeeze(), "labels": torch.tensor(label)}
+     
+dataset = ClinicalDataset(texts=train_texts, labels=train_labels, tokenizer=tokenizer)
 
+# Split dataset into training and validation sets
+train_size = int(0.8 * len(dataset))
+val_size = len(dataset) - train_size
+train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
 
-dataset = ClinicalDataset(texts=train_texts, labels=train_labels, tokenizer=tokenizer)
-# Fine-tune the pre-trained model on your clinical dataset
-from transformers import Trainer, TrainingArguments
-    
+# Fine-tune pre-trained model on clinical dataset
 training_args = TrainingArguments(
          output_dir='./results',          # output directory
          num_train_epochs=3,              # total number of training epochs
@@ -43,12 +47,11 @@ training_args = TrainingArguments(
          warmup_steps=500,                # number of warmup steps for learning rate scheduler
          weight_decay=0.01,               # strength of weight decay
          logging_dir='./logs',            # directory for storing logs
-         logging_steps=10, )
-    
+         logging_steps=10,)
 trainer = Trainer(
       model=model,
       args=training_args,
-      train_dataset=dataset,
+      train_dataset=train_dataset,
       eval_dataset=val_dataset,
       data_collator=lambda data: {'input_ids': torch.stack([f['input_ids'] for f in data]),
                                 'attention_mask': torch.stack([f['attention_mask'] for f in data]),