Update app.py

app.py

@@ -12,7 +12,7 @@ import requests
 
 app = FastAPI()
 
-# ✅ Correct Hugging Face Dataset URL
+# ✅ Correct dataset URL from Hugging Face (Replace if needed)
 DATASET_URL = "https://huggingface.co/datasets/SailajaS/CDART/resolve/main/train.csv"
 
 # File path for saving dataset
@@ -32,29 +32,35 @@ def download_dataset():
             print(f"❌ Failed to download dataset: {response.status_code}")
             raise Exception("Dataset download failed.")
 
-# Download dataset at startup
+# ✅ Download dataset at startup
 download_dataset()
 
-# Load dataset
+# ✅ Load dataset
 df = pd.read_csv(DATASET_PATH)
 
-# Encode categorical variables
+# ✅ Check if necessary columns exist
+required_columns = ["Case Problem", "Feedback"]
+for col in required_columns:
+    if col not in df.columns:
+        raise Exception(f"❌ Column '{col}' is missing from the dataset!")
+
+# ✅ Encode categorical variables
 encoder = LabelEncoder()
 df["Case Problem"] = encoder.fit_transform(df["Case Problem"])
 df["Feedback"] = encoder.fit_transform(df["Feedback"])
 
-# Train Model
+# ✅ Train Model
 X = df[["Case Problem"]]
 y = df["Feedback"]
 X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
 model = RandomForestClassifier(n_estimators=100, random_state=42)
 model.fit(X_train, y_train)
 
-# Save model
+# ✅ Save model
 joblib.dump(model, "feedback_model.pkl")
 print("✅ Model trained successfully!")
 
-# API Input Model
+# ✅ API Input Model
 class PredictionInput(BaseModel):
     case_problem: str
 
@@ -64,22 +70,28 @@ def predict_feedback(data: PredictionInput):
     if model is None:
         return {"error": "Model is not trained yet."}
 
-    case_problem_encoded = encoder.transform([data.case_problem])
-    prediction = model.predict([[case_problem_encoded[0]]])
-    feedback_predicted = encoder.inverse_transform(prediction)[0]
-    return {"Predicted Feedback": feedback_predicted}
+    try:
+        case_problem_encoded = encoder.transform([data.case_problem])
+        prediction = model.predict([[case_problem_encoded[0]]])
+        feedback_predicted = encoder.inverse_transform(prediction)[0]
+        return {"Predicted Feedback": feedback_predicted}
+    except Exception as e:
+        return {"error": str(e)}
 
-# Gradio UI
+# ✅ Gradio UI
 def gradio_interface(case_problem):
     if model is None:
         return "Model not trained yet."
 
-    case_problem_encoded = encoder.transform([case_problem])
-    prediction = model.predict([[case_problem_encoded[0]]])
-    feedback_predicted = encoder.inverse_transform(prediction)[0]
-    return f"Predicted Feedback: {feedback_predicted}"
+    try:
+        case_problem_encoded = encoder.transform([case_problem])
+        prediction = model.predict([[case_problem_encoded[0]]])
+        feedback_predicted = encoder.inverse_transform(prediction)[0]
+        return f"Predicted Feedback: {feedback_predicted}"
+    except:
+        return "Invalid case problem. Please enter a valid category."
 
-# Start both API & Gradio
+# ✅ Start both API & Gradio
 def start_app():
     """ Start API and Gradio Interface """
     gr_interface = gr.Interface(fn=gradio_interface, inputs="text", outputs="text")