import gradio as gr from transformers import pipeline import torch import os from typing import Union, List # ============================== # ⚙️ Configuration # ============================== MODEL_NAME = os.getenv("MODEL_NAME", "MoritzLaurer/deberta-v3-base-zeroshot-v2.0") # Optional: specify cache dir HF_HOME = os.getenv("HF_HOME", "/tmp/huggingface") os.environ["TRANSFORMERS_CACHE"] = HF_HOME DEVICE = 0 if torch.cuda.is_available() else -1 # ============================== # 🔧 Load Model Pipeline # ============================== try: print("🚀 Loading zero-shot classification pipeline...") severity_classifier = pipeline( "zero-shot-classification", model=MODEL_NAME, device=DEVICE ) print(f"✅ Model loaded successfully: {MODEL_NAME}") except Exception as e: raise RuntimeError(f"❌ Failed to load model: {e}") # ============================== # 🧠 Classification Function # ============================== def classify_severity(side_effects_input: Union[str, List[str]]): """ Accepts a string (comma-separated) or a list of side effects. Returns a dictionary mapping side effects -> severity. """ try: # Handle both list and string input types if isinstance(side_effects_input, str): side_effects_list = [ se.strip() for se in side_effects_input.split(",") if se.strip() ] elif isinstance(side_effects_input, list): side_effects_list = [se.strip() for se in side_effects_input if se.strip()] else: return {"error": "Invalid input format. Expected string or list."} if not side_effects_list: return {} candidate_labels = ["Mild", "Moderate", "Severe", "Life-threatening", "Death"] # Perform classification results = severity_classifier( sequences=side_effects_list, candidate_labels=candidate_labels, multi_label=False ) # Normalize single result into list if isinstance(results, dict): results = [results] # Convert to dict { side_effect: severity } output_dict = { result["sequence"]: result["labels"][0].lower() for result in results } return output_dict except Exception as e: return {"error": str(e)} # ============================== # 🎨 Gradio Interface # ============================== demo = gr.Interface( fn=classify_severity, inputs=[ gr.Textbox( label="Side Effects", placeholder="Enter side effects, separated by commas (e.g., headache, nausea, dizziness)" ) ], outputs=gr.JSON(label="Severity Classification"), title="⚕️ Side Effect Severity Classifier", description="Classifies side effect severity using a DeBERTa/BART model for zero-shot classification.", examples=[ ["headache, fatigue"], ["chest pain, difficulty breathing"], ["slight rash"] ] ) # ============================== # 🚀 Launch App # ============================== if __name__ == "__main__": demo.launch(server_name="0.0.0.0", server_port=7860) # import gradio as gr # from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer # from pathlib import Path # import torch # import os # # --- 1. Configuration --- # MODEL_NAME = "MoritzLaurer/deberta-v3-large-zeroshot-v2.0" # LOCAL_MODEL_DIR = Path(__file__).parent / "models--MoritzLaurer--deberta-v3-large-zeroshot-v2.0" / "snapshots" / "cf44676c28ba7312e5c5f8f8d2c22b3e0c9cdae2" # DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # # --- 2. Model Loading with Fallback --- # try: # # First try loading from local directory # if LOCAL_MODEL_DIR.exists(): # print("🔄 Attempting to load model from local directory...") # model = AutoModelForSequenceClassification.from_pretrained(str(LOCAL_MODEL_DIR)) # tokenizer = AutoTokenizer.from_pretrained(str(LOCAL_MODEL_DIR)) # severity_classifier = pipeline( # "zero-shot-classification", # model=model, # tokenizer=tokenizer, # device=DEVICE # ) # print(f"✅ Successfully loaded model from local path: {LOCAL_MODEL_DIR}") # else: # # Fallback to downloading from HuggingFace Hub # print("🌐 Local model not found, downloading from HuggingFace Hub...") # severity_classifier = pipeline( # "zero-shot-classification", # model=MODEL_NAME, # device=DEVICE # ) # print(f"✅ Successfully downloaded and loaded model: {MODEL_NAME}") # except Exception as e: # # Detailed error diagnostics # if LOCAL_MODEL_DIR.exists(): # missing_files = [ # f for f in ["config.json", "model.safetensors", "special_tokens_map.json", # "tokenizer_config.json", "spm.model"] # if not (LOCAL_MODEL_DIR / f).exists() # ] # if missing_files: # raise RuntimeError(f"❌ Missing required model files: {missing_files}") # raise RuntimeError(f"❌ Failed to load model. Error: {str(e)}") # # --- 3. Classification Function --- # def classify_severity(side_effects_str: str): # if not side_effects_str or not side_effects_str.strip(): # return {} # side_effects_list = [se.strip() for se in side_effects_str.split(',')] # candidate_labels = [ # "Mild", # "Moderate", # "Severe", # "Life-threatening", # "Death" # ] # results = severity_classifier( # sequences=side_effects_list, # candidate_labels=candidate_labels, # multi_label=False # ) # return {result['sequence']: result['labels'][0].lower() for result in results} # # --- 4. Gradio Interface --- # demo = gr.Interface( # fn=classify_severity, # inputs=gr.Textbox( # label="Side Effects", # placeholder="Enter side effects, separated by commas (e.g., headache, nausea, dizziness)" # ), # outputs=gr.JSON(label="Severity Classification"), # title="⚕️ Side Effect Severity Classifier", # description="Classifies side effect severity using DeBERTa model", # examples=[ # ["headache, fatigue"], # ["chest pain, difficulty breathing"], # ["slight rash"] # ] # ) # # --- 5. Launch --- # if __name__ == "__main__": # demo.launch(server_name="0.0.0.0", server_port=7860) # import gradio as gr # from transformers import pipeline # from pathlib import Path # Import Path for handling file paths # import torch # # --- 1. Define the LOCAL PATH to the model --- # # This path exactly matches the folder name in your repository screenshot. # MODEL_PATH = Path(__file__).parent / "models--MoritzLaurer--deberta-v3-large-zeroshot-v2.0" / "snapshots" / "cf44676c28ba7312e5c5f8f8d2c22b3e0c9cdae2" # DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # # --- 2. Load the Zero-Shot Classification model from the LOCAL PATH --- # try: # # We now pass the local Path object to the model argument. # severity_classifier = pipeline( # "zero-shot-classification", # model=MODEL_PATH, # device=DEVICE # ) # print(f"✅ Successfully loaded model from local path: '{MODEL_PATH}' on {DEVICE}") # except Exception as e: # # If this fails, it's because the files inside the folder are missing or incorrect. # raise RuntimeError(f"❌ Failed to load model from local path. Ensure the directory '{MODEL_PATH}' exists and contains all necessary files (config, model, tokenizer). Error: {str(e)}") # # --- 3. Implement your specific classification logic (optimized version) --- # def classify_severity(side_effects_str: str): # """ # Takes a comma-separated string of side effects, classifies them all # at once (batching), and returns a dictionary of their severities. # """ # # Return an empty dictionary if the input is empty or just whitespace # if not side_effects_str or not side_effects_str.strip(): # return {} # # Split the input string into a list of individual side effects # side_effects_list = [se.strip() for se in side_effects_str.split(',')] # # These are the fixed labels for our specific severity classification task # candidate_labels = [ # "Mild", # "Moderate", # "Severe", # "Life-threatening", # "Death" # ] # # --- PERFORMANCE OPTIMIZATION --- # # We call the classifier only ONCE with the entire list of side effects. # # This is much faster than calling it in a for-loop for each one. # results = severity_classifier( # sequences=side_effects_list, # candidate_labels=candidate_labels, # multi_label=False # We want only the single best label per side effect # ) # # Process the results into a clean dictionary, making the label lowercase like your original code. # severity_dict = { # result['sequence']: result['labels'][0].lower() for result in results # } # return severity_dict # # --- 4. Create the Gradio Interface (which also creates the API endpoint) --- # demo = gr.Interface( # fn=classify_severity, # inputs=gr.Textbox( # label="Side Effects", # placeholder="Enter side effects, separated by commas (e.g., headache, nausea, dizziness)" # ), # outputs=gr.JSON(label="Severity Classification"), # title="⚕️ Side Effect Severity Classifier (Local Model)", # description=f"Enter one or more side effects to classify their severity. The model is loaded from the repository.", # examples=[ # ["headache, fatigue"], # ["chest pain, difficulty breathing"], # ["slight rash"] # ] # ) # # --- 5. Launch the app --- # # This makes the web UI and the background API available. # demo.launch() # import gradio as gr # from transformers import pipeline # import torch # # --- 1. Define the Model to be loaded from Hugging Face Hub --- # MODEL_NAME = "MoritzLaurer/deberta-v3-large-zeroshot-v2.0" # DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # # --- 2. Load the Zero-Shot Classification model --- # try: # # We specify the task and the model's public name. # # The library handles downloading/caching automatically. # severity_classifier = pipeline( # "zero-shot-classification", # model=MODEL_NAME, # device=DEVICE # ) # print(f"✅ Successfully loaded model '{MODEL_NAME}' on {DEVICE}") # except Exception as e: # raise RuntimeError(f"❌ Failed to load the classification model: {str(e)}") # # --- 3. Implement your specific classification logic --- # # This function is optimized to be much faster than a for-loop. # def classify_severity(side_effects_str: str): # """ # Takes a comma-separated string of side effects, classifies them all # at once, and returns a dictionary of their severities. # """ # # Return an empty dictionary if the input is empty # if not side_effects_str or not side_effects_str.strip(): # return {} # # Split the input string into a list of individual side effects # side_effects_list = [se.strip() for se in side_effects_str.split(',')] # # These are the fixed labels for our specific task # candidate_labels = [ # "Mild", # "Moderate", # "Severe", # "Life-threatening", # "Death" # ] # # --- PERFORMANCE OPTIMIZATION --- # # We call the classifier only ONCE with the entire list (batching). # results = severity_classifier( # sequences=side_effects_list, # candidate_labels=candidate_labels, # multi_label=False # We want only the single best label for each side effect # ) # # Process the results into a clean dictionary # severity_dict = { # result['sequence']: result['labels'][0] for result in results # } # return severity_dict # # --- 4. Create the Gradio Interface (which also creates the API endpoint) --- # demo = gr.Interface( # fn=classify_severity, # inputs=gr.Textbox( # label="Side Effects", # placeholder="Enter side effects, separated by commas (e.g., headache, nausea, dizziness)" # ), # outputs=gr.JSON(label="Severity Classification"), # title="⚕️ Side Effect Severity Classifier", # description=f"Enter one or more side effects to classify their severity using the {MODEL_NAME} model.", # examples=[ # ["headache, fatigue"], # ["chest pain, difficulty breathing"], # ["slight rash"] # ] # ) # # --- 5. Launch the app --- # # This makes the web UI and the background API available. # demo.launch() # import gradio as gr # from transformers import pipeline # import torch # # --- 1. Load Model from Hugging Face Hub --- # # We specify the model by its name. 'transformers' will download it. # MODEL_NAME = "MoritzLaurer/deberta-v3-large-zeroshot-v2.0" # DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # try: # # --- 2. Change the pipeline task to "zero-shot-classification" --- # classifier = pipeline( # "zero-shot-classification", # model=MODEL_NAME, # device=DEVICE # ) # print(f"✅ Model '{MODEL_NAME}' loaded successfully on {DEVICE}") # except Exception as e: # raise RuntimeError(f"❌ Model loading failed: {e}") # # --- 3. Update the function to accept labels --- # # The function now needs the text AND the labels to classify against. # def analyze(text, candidate_labels): # # The user will input labels as a comma-separated string, e.g., "positive, negative, neutral" # # We need to turn this into a list of strings. # labels = [label.strip() for label in candidate_labels.split(',')] # # The zero-shot classifier returns a dictionary with labels and scores # result = classifier(text, candidate_labels=labels) # return result # # --- 4. Update the Gradio Interface for two inputs --- # demo = gr.Interface( # fn=analyze, # # We now have two inputs: the text and a textbox for the labels # inputs=[ # gr.Textbox(label="Input Text", placeholder="Enter some text here..."), # gr.Textbox(label="Candidate Labels", placeholder="e.g., positive, negative, neutral") # ], # # The output is a JSON object, which is perfect for the result # outputs=gr.JSON(label="Classification Result"), # title="🔍 Zero-Shot Text Classification", # description="Enter text and a comma-separated list of categories to classify it into. Model: MoritzLaurer/deberta-v3-large-zeroshot-v2.0", # # Examples must now provide values for both inputs # examples=[ # ["The new AI regulations will have a huge impact.", "politics, business, technology"], # ["I am so excited for the game tonight!", "sports, entertainment, work"], # ["This movie was fantastic, I loved the acting.", "positive, negative"] # ] # ) # # Launch the app # demo.launch()