Update app.py

app.py

@@ -128,17 +128,13 @@ def classify_website(url):
     try:
         # Prepare the input prompt for the model
         prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
-
 ### Instruction:
 Categorize the website into one of the 3 categories:
-
 1) OTHER
 2) NEWS/BLOG
 3) E-commerce
-
 ### Input:
 {translated}
-
 ### Response:"""
 
         inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
@@ -159,13 +155,35 @@ Categorize the website into one of the 3 categories:
         logging.exception(e)
         return str(e)
 
+
+def classify_urls_from_csv(csv_file):
+    # Read CSV file and extract URLs from the first column
+    df = pd.read_csv(csv_file)
+    urls = df.iloc[:, 0].tolist()
+
+    # Classify each URL and store the results
+    predictions = []
+    for url in urls:
+        prediction = classify_website(url)
+        predictions.append(prediction)
+
+    # Add predictions as a new column in the dataframe
+    df['Prediction'] = predictions
+
+    # Save the results to a new CSV file
+    output_file = "predictions.csv"
+    df.to_csv(output_file, index=False)
+
+    return output_file
+
+
 # Create a Gradio interface
 iface = gr.Interface(
-    fn=classify_website,
-    inputs="text",
-    outputs="text",
+    fn=classify_urls_from_csv,
+    inputs=gr.File(label="Upload CSV (First column must contain URLs)"),
+    outputs=gr.File(label="Download CSV with Predictions"),
     title="Website Categorization",
-    description="Categorize a website into one of the 3 categories: OTHER, NEWS/BLOG, or E-commerce."
+    description="Upload a CSV file with URLs in the first column and categorize each website into one of the 3 categories: OTHER, NEWS/BLOG, or E-commerce."
 )
 
 # Launch the interface