Update app.py

app.py

@@ -27,27 +27,23 @@ dataset_path = hf_hub_download(repo_id="Fredaaaaaa/hybrid_model", filename="labe
 df = pd.read_csv(dataset_path, encoding='ISO-8859-1')
 print(f"Dataset loaded successfully! Shape: {df.shape}")
 
-# Create a set of all unique drugs in the dataset for validation
-all_drugs = set()
+# Check the columns and display first few rows for debugging
+print(df.columns)
+print(df.head())
 
-# Check which columns contain drug names
-drug_columns = []
-for col in df.columns:
-    if 'drug' in col.lower() or 'medication' in col.lower():
-        drug_columns.append(col)
-        # Add all drugs from this column to our set after cleaning
-        clean_drugs = df[col].dropna().astype(str).apply(lambda x: x.strip().lower())
-        all_drugs.update(clean_drugs.unique())
+# Get unique severity classes from the dataset
+unique_classes = df['severity'].unique()
+print(f"Unique severity classes in dataset: {unique_classes}")
 
 # Calculate class weights to handle imbalanced classes
-
-
-# Correct the 'classes' parameter to be a numpy.ndarray
-class_weights = compute_class_weight('balanced', classes=np.array([0, 1, 2, 3]), y=df['severity'])
-
+# Use the unique classes from the dataset for the `classes` parameter
+class_weights = compute_class_weight('balanced', classes=np.unique(unique_classes), y=df['severity'])
 class_weights = torch.tensor(class_weights, dtype=torch.float).to(device)
 loss_fn = torch.nn.CrossEntropyLoss(weight=class_weights)
 
+# The rest of your code follows here...
+
+
 # Function to properly clean drug names
 def clean_drug_name(drug_name):
     if not drug_name: