Update app to use internal data

app.py

@@ -1,19 +1,27 @@
 import streamlit as st
 import torch
 from torch.utils.data import Dataset, random_split
-from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments, default_data_collator
 
-# 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):")
+# Generate sample clinical text and labels
+sample_data = [
+    ("Had successful surgery today. Feeling relieved.", "surgery"),
+    ("Started new medication for pain management.", "non-surgery"),
+    ("Scheduled for surgery next week. Nervous but hopeful.", "surgery"),
+    ("Attended a seminar on non-surgical treatments.", "non-surgery"),
+]
 
-# Convert comma-separated values into lists
-train_texts = train_texts.split(",")
-train_labels = train_labels.split(",")
+train_texts, train_labels = zip(*sample_data)
+
+# Logging and Outputs
+st.write("Sample data:")
+for text, label in zip(train_texts, train_labels):
+    st.write(f"Text: {text}\nLabel: {label}\n")
 
 # Load pre-trained model and tokenizer
-model = AutoModelForCausalLM.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
-tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
+model_name = "distilbert-base-uncased"  # You can use any suitable classification model
+model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 # Create PyTorch Dataset object
 class ClinicalDataset(Dataset):
@@ -31,6 +39,11 @@ class ClinicalDataset(Dataset):
              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)}
      
+
+
+# Data Collator
+data_collator = default_data_collator
+
 dataset = ClinicalDataset(texts=train_texts, labels=train_labels, tokenizer=tokenizer)
 
 # Split dataset into training and validation sets
@@ -48,6 +61,7 @@ training_args = TrainingArguments(
          weight_decay=0.01,               # strength of weight decay
          logging_dir='./logs',            # directory for storing logs
          logging_steps=10,)
+
 trainer = Trainer(
       model=model,
       args=training_args,
@@ -56,4 +70,13 @@ trainer = Trainer(
       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]),
                                  'labels': torch.stack([f['labels'] for f in data])}, )
+
+
+st.write("Training started...")
 trainer.train()
+st.write("Training completed.")
+
+# Logging Training Output
+st.write("Training logs:")
+with open('./logs/train.log', 'r') as log_file:
+    st.code(log_file.read())