Spaces:
Running
Running
File size: 3,041 Bytes
2bc3020 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoModel
import dask.dataframe as dd
from datasets import load_dataset
import torch
# Load models and tokenizer
def load_models():
# Load model 1
model_1 = AutoModel.from_pretrained("Canstralian/RedTeamAI")
# Load model 2
model_2 = AutoModel.from_pretrained("mradermacher/BashCopilot-6B-preview-GGUF")
# Load tokenizer and sequence classification model
tokenizer = AutoTokenizer.from_pretrained("bash1130/bert-base-finetuned-ynat")
model_3 = AutoModelForSequenceClassification.from_pretrained("bash1130/bert-base-finetuned-ynat")
return model_1, model_2, tokenizer, model_3
# Load dataset using Dask
def load_data():
# Example of loading a dataset using Dask (adjust paths as necessary)
splits = {'creative_content': 'data/creative_content-00000-of-00001.parquet'}
df = dd.read_parquet("hf://datasets/microsoft/orca-agentinstruct-1M-v1/" + splits["creative_content"])
return df.head()
# Function for model inference
def infer_model(input_text, model_type):
# Choose the model based on the input (you can add more models or conditions as needed)
if model_type == 'RedTeamAI':
model = models[0]
elif model_type == 'BashCopilot':
model = models[1]
elif model_type == 'BertModel':
model = models[3]
inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
outputs = model(**inputs)
return outputs.logits.argmax(dim=-1).item()
else:
return "Model type not recognized."
# If you need to generate outputs based on the models directly, you can use:
# outputs = model.generate(input_text) or other inference methods depending on the model.
return f"Model {model_type} inference not implemented yet."
# Gradio Interface setup
def build_interface():
# Load models and data
model_1, model_2, tokenizer, model_3 = load_models()
global models
models = [model_1, model_2, tokenizer, model_3]
# Load the dataset (example function, you can add more functionality)
data_preview = load_data()
print(f"Dataset preview: {data_preview}")
# Create Gradio interface
with gr.Blocks() as demo:
gr.Markdown("# Chagrin AI - Model Inference & Dataset Explorer")
# Model selection dropdown
model_type = gr.Dropdown(choices=["RedTeamAI", "BashCopilot", "BertModel"], label="Choose Model")
# Textbox for user input
input_text = gr.Textbox(label="Enter your input text")
# Button to trigger inference
result = gr.Textbox(label="Inference Result")
submit_btn = gr.Button("Run Inference")
submit_btn.click(infer_model, inputs=[input_text, model_type], outputs=result)
demo.launch()
# Run the app
if __name__ == "__main__":
build_interface()
|