Update app.py

app.py

@@ -15,10 +15,15 @@ import gradio as gr
 import openai
 from datasets import load_dataset
 
-# Ensure Hugging Face CLI is authenticated
-if "HF_TOKEN" not in os.environ:
-    print("Please authenticate with Hugging Face CLI or set HF_TOKEN as an environment variable.")
-    exit(1)
+# Ensure Hugging Face token exists
+HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    raise ValueError("Hugging Face token is not set. Please set HF_TOKEN as an environment variable.")
+
+# Ensure OpenAI API key exists
+openai.api_key = os.getenv("OPENAI_API_KEY")
+if not openai.api_key:
+    raise ValueError("OpenAI API key is not set. Please set OPENAI_API_KEY as an environment variable.")
 
 # Explicitly define dataset file paths
 data_files = {
@@ -33,7 +38,6 @@ try:
     train_data = dataset["train"].to_pandas()
     validation_data = dataset["validation"].to_pandas()
     print("Dataset loaded successfully.")
-    print("Train dataset columns:", train_data.columns)
 except Exception as e:
     print(f"Failed to load dataset: {e}")
     train_data = pd.DataFrame()  # Fallback to empty DataFrame
@@ -56,11 +60,6 @@ for column in required_columns:
 tokenized_train = [doc.split() for doc in train_data["text"]]
 bm25 = BM25Okapi(tokenized_train)
 
-# Set OpenAI API key
-openai.api_key = os.getenv("OPENAI_API_KEY")
-if not openai.api_key:
-    raise ValueError("OpenAI API key is not set. Please set it as an environment variable.")
-
 # Few-shot classification function
 def classify_text(input_text, k=20):
     # Tokenize input text
@@ -102,9 +101,16 @@ def classify_text(input_text, k=20):
             model="gpt-4",
             temperature=0,
         )
+        print("OpenAI Response:", response)
+        if "choices" not in response or not response["choices"]:
+            raise ValueError("Invalid or empty response from OpenAI.")
         results = response.choices[0].message.content
         return json.loads(results)
+    except json.JSONDecodeError:
+        print("Failed to decode JSON from OpenAI response.")
+        return {"error": "Failed to decode JSON from OpenAI response."}
     except Exception as e:
+        print(f"Error occurred: {e}")
         return {"error": str(e)}
 
 # Gradio Interface