Update app.py

app.py

@@ -1,39 +1,67 @@
-import gradio as gr
-import tensorflow as tf
+import os
 import requests
-import tempfile
-
-# Define the URL for your model and dataset
-model_url = 'https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/bacterial_morphology_classification_model.h5'
-
-# Download the model from the URL and load it into a temporary file
-response = requests.get(model_url, stream=True)
-if response.status_code == 200:
-    with tempfile.NamedTemporaryFile(delete=False, suffix='.h5') as temp_file:
-        temp_file.write(response.content)
-        temp_file_path = temp_file.name
-        model = tf.keras.models.load_model(temp_file_path)
-else:
-    raise Exception(f"Failed to download the model. Status code: {response.status_code}")
-
-# Define the function for prediction
-def predict_image(img):
-    img = img.resize((224, 224))
-    img_array = tf.keras.preprocessing.image.img_to_array(img) / 255.0
-    img_array = tf.expand_dims(img_array, axis=0)
-    prediction = model.predict(img_array)
-    classes = ['cocci', 'bacilli', 'spirilla']
-    predicted_class = classes[prediction.argmax()]
-    return predicted_class
-
-# Create the Gradio interface
-iface = gr.Interface(
-    fn=predict_image,
-    inputs=gr.inputs.Image(shape=(224, 224)),
-    outputs=gr.outputs.Label(num_top_classes=3),
-    title="Bacterial Morphology Classification",
-    description="Upload an image of bacterial morphology to classify it as cocci, bacilli, or spirilla."
+from datasets import load_dataset
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import torch
+
+# Step 1: Set up environment and check paths
+MODEL_PATH = "yolac/BacterialMorphologyClassification"
+DATASET_PATH = "yolac/BacterialMorphologyClassification"
+
+# Step 2: Load the model
+try:
+    model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH, num_labels=3)
+    print("Model loaded successfully.")
+except Exception as e:
+    print(f"Failed to load the model. Error: {e}")
+
+# Step 3: Load the tokenizer (if needed)
+try:
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+    print("Tokenizer loaded successfully.")
+except Exception as e:
+    print(f"Failed to load the tokenizer. Error: {e}")
+
+# Step 4: Load the dataset
+try:
+    dataset = load_dataset(DATASET_PATH, split="train")
+    print("Dataset loaded successfully.")
+except Exception as e:
+    print(f"Failed to load the dataset. Error: {e}")
+
+# Step 5: Preprocess and prepare data for model input (example code)
+def preprocess_data(example):
+    # Add any necessary preprocessing steps here, e.g., tokenization
+    return tokenizer(example['text'], padding="max_length", truncation=True)
+
+# Apply preprocessing
+dataset = dataset.map(preprocess_data, batched=True)
+
+# Step 6: Set up training arguments (use the `Trainer` class if needed)
+from transformers import Trainer, TrainingArguments
+
+training_args = TrainingArguments(
+    output_dir="./results",
+    evaluation_strategy="epoch",
+    learning_rate=2e-5,
+    per_device_train_batch_size=8,
+    per_device_eval_batch_size=8,
+    num_train_epochs=3,
+    weight_decay=0.01,
+    logging_dir="./logs",
+)
+
+# Initialize the Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=dataset,
+    tokenizer=tokenizer,
 )
 
-# Launch the Gradio app
-iface.launch()
+# Step 7: Train the model
+try:
+    trainer.train()
+    print("Training completed successfully.")
+except Exception as e:
+    print(f"Failed to train the model. Error: {e}")