Spaces:
Sleeping
Sleeping
| import pandas as pd | |
| import numpy as np | |
| import torch | |
| from sklearn.preprocessing import LabelEncoder | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
| import gradio as gr | |
| import gspread | |
| from oauth2client.service_account import ServiceAccountCredentials | |
| import os | |
| import json | |
| # ========================= | |
| # PASSWORD (from env) | |
| # ========================= | |
| APP_PASSWORD = os.environ.get("APP_PASSWORD") | |
| if APP_PASSWORD is None: | |
| raise ValueError("APP_PASSWORD environment variable not set") | |
| # ========================= | |
| # THEMES | |
| # ========================= | |
| THEMES = { | |
| "Lifestyle": "Lifestyle: everyday life, hobbies, routines, hanging out, personal moments", | |
| "Travel & Experience": "Travel & Experience: trips, visiting places, tourism, events, sports games, concerts", | |
| "Reflection": "Reflection: deep thoughts, life lessons, gratitude, personal growth, emotional reflection", | |
| "Brand & Business": "Brand & Business: marketing, promoting a brand, entrepreneurship, product announcements, advertisements", | |
| "Education & Value": "Education & Value: advice, tips, teaching something, sharing knowledge", | |
| "Relationships": "Relationships: family, friends, romantic, social connections", | |
| "Performance & Sport": "Performance & Sport: competitions, sports achievements, athletic performance" | |
| } | |
| label_encoder = LabelEncoder() | |
| label_encoder.fit(list(THEMES.keys())) | |
| # ========================= | |
| # GOOGLE SHEETS AUTH | |
| # ========================= | |
| scope = [ | |
| "https://spreadsheets.google.com/feeds", | |
| "https://www.googleapis.com/auth/drive" | |
| ] | |
| sheet = None | |
| try: | |
| creds_json_str = os.environ.get("GDRIVE_SERVICE_ACCOUNT_JSON") | |
| if not creds_json_str: | |
| print("CRITICAL ERROR: Missing GDRIVE_SERVICE_ACCOUNT_JSON") | |
| else: | |
| creds_info = json.loads(creds_json_str) | |
| creds = ServiceAccountCredentials.from_json_keyfile_dict(creds_info, scope) | |
| client = gspread.authorize(creds) | |
| sheet = client.open("ClassifierTraining").get_worksheet(0) | |
| print(f"Connected to sheet: {sheet.title}") | |
| except Exception as e: | |
| print(f"Auth Failure: {e}") | |
| # ========================= | |
| # MODEL LOAD | |
| # ========================= | |
| model_name = "nirvairkataria/caption-classifier-model" | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| model = AutoModelForSequenceClassification.from_pretrained(model_name) | |
| model.to("cuda" if torch.cuda.is_available() else "cpu") | |
| model.eval() | |
| # ========================= | |
| # FUNCTIONS | |
| # ========================= | |
| def classify_caption_trained_model(caption): | |
| inputs = tokenizer(caption, return_tensors="pt", truncation=True, padding=True) | |
| inputs = {k: v.to(model.device) for k, v in inputs.items()} | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| logits = outputs.logits | |
| probs = torch.softmax(logits, dim=1).cpu().numpy()[0] | |
| label_idx = np.argmax(probs) | |
| pred_label = label_encoder.inverse_transform([label_idx])[0] | |
| confidence = round(float(probs[label_idx]), 2) | |
| all_themes = list(THEMES.keys()) | |
| return pred_label, confidence, all_themes | |
| def save_feedback_to_sheet(caption, corrected_theme): | |
| if sheet is None: | |
| return "Error: Google Sheet authentication failed." | |
| try: | |
| sheet.append_row([caption, corrected_theme]) | |
| return "Feedback submitted successfully! It will be used when the model is retrained in order to improve accuracy for future captions." | |
| except Exception as e: | |
| return f"Failed: {e}" | |
| def process_csv(file_obj): | |
| df = pd.read_csv(file_obj.name) | |
| predicted_themes, confidences = [], [] | |
| for caption in df["message"]: | |
| theme, conf, _ = classify_caption_trained_model(caption) | |
| predicted_themes.append(theme) | |
| confidences.append(conf) | |
| df["Predicted Theme"] = predicted_themes | |
| df["Confidence"] = confidences | |
| output_file = "classified_results.csv" | |
| df.to_csv(output_file, index=False) | |
| return df, output_file | |
| def handle_classification(caption): | |
| pred_label, confidence, all_themes = classify_caption_trained_model(caption) | |
| return pred_label, confidence, gr.update(choices=all_themes, value=pred_label) | |
| # ========================= | |
| # PASSWORD CHECK | |
| # ========================= | |
| def check_password(pw): | |
| if pw == APP_PASSWORD: | |
| return gr.update(visible=False), gr.update(visible=True), "" | |
| else: | |
| return gr.update(), gr.update(), "Incorrect password" | |
| # ========================= | |
| # UI | |
| # ========================= | |
| with gr.Blocks() as demo: | |
| # LOGIN SCREEN | |
| with gr.Column(visible=True) as login_screen: | |
| gr.Markdown("## 🔒 Enter Password") | |
| password_input = gr.Textbox(type="password", label="Password") | |
| login_btn = gr.Button("Enter") | |
| login_status = gr.Textbox(label="", interactive=False) | |
| # MAIN APP (HIDDEN) | |
| with gr.Column(visible=False) as app_screen: | |
| gr.Markdown("# Caption Classifier") | |
| with gr.Tab("Single Caption"): | |
| caption_input = gr.Textbox(label="Enter Caption", lines=4) | |
| with gr.Row(): | |
| theme_output = gr.Textbox(label="Predicted Theme", interactive=False) | |
| confidence_output = gr.Textbox(label="Confidence", interactive=False) | |
| classify_btn = gr.Button("Classify") | |
| corrected_theme_dropdown = gr.Dropdown(label="Corrected Theme", choices=list(THEMES.keys())) | |
| submit_btn = gr.Button("Submit Correction") | |
| status_output = gr.Textbox(label="Status", interactive=False) | |
| classify_btn.click( | |
| handle_classification, | |
| inputs=caption_input, | |
| outputs=[theme_output, confidence_output, corrected_theme_dropdown] | |
| ) | |
| submit_btn.click( | |
| save_feedback_to_sheet, | |
| inputs=[caption_input, corrected_theme_dropdown], | |
| outputs=status_output | |
| ) | |
| with gr.Tab("CSV Upload"): | |
| file_input = gr.File(label="Upload CSV") | |
| output_table = gr.Dataframe() | |
| download_file = gr.File() | |
| process_btn = gr.Button("Run") | |
| process_btn.click( | |
| process_csv, | |
| inputs=file_input, | |
| outputs=[output_table, download_file] | |
| ) | |
| # LOGIN ACTION | |
| login_btn.click( | |
| check_password, | |
| inputs=password_input, | |
| outputs=[login_screen, app_screen, login_status] | |
| ) | |
| # ========================= | |
| # LAUNCH | |
| # ========================= | |
| demo.launch(server_name="0.0.0.0", server_port=7860) |