Create app.py

app.py

@@ -0,0 +1,168 @@
+import streamlit as st
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel, PretrainedConfig, AutoTokenizer, AutoModelForSequenceClassification, AutoModelForSeq2SeqLM
+from huggingface_hub import login
+import os
+import time
+
+# Model Architecture
+class TinyTransformer(nn.Module):
+    def __init__(self, vocab_size, embed_dim, num_heads, ff_dim, num_layers):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        self.pos_encoding = nn.Parameter(torch.zeros(1, 512, embed_dim))
+        encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, dim_feedforward=ff_dim, batch_first=True)
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+        self.fc = nn.Linear(embed_dim, 1)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        x = self.embedding(x) + self.pos_encoding[:, :x.size(1), :]
+        x = self.transformer(x)
+        x = x.mean(dim=1)  # Global average pooling
+        x = self.fc(x)
+        return self.sigmoid(x)
+
+class TinyTransformerConfig(PretrainedConfig):
+    model_type = "tiny_transformer"
+
+    def __init__(
+        self,
+        vocab_size=30522,
+        embed_dim=64,
+        num_heads=2,
+        ff_dim=128,
+        num_layers=4,
+        max_position_embeddings=512,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.ff_dim = ff_dim
+        self.num_layers = num_layers
+        self.max_position_embeddings = max_position_embeddings
+
+class TinyTransformerForSequenceClassification(PreTrainedModel):
+    config_class = TinyTransformerConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = 1
+        self.transformer = TinyTransformer(
+            config.vocab_size,
+            config.embed_dim,
+            config.num_heads,
+            config.ff_dim,
+            config.num_layers
+        )
+
+    def forward(self, input_ids, attention_mask=None):
+        outputs = self.transformer(input_ids)
+        return {"logits": outputs}
+
+# Load models and tokenizers
+@st.cache_resource
+def load_models_and_tokenizers(hf_token):
+    login(token=hf_token)
+    device = torch.device("cpu")  # forcing CPU as overhead of inference on GPU slows down the inference
+    
+    models = {}
+    tokenizers = {}
+    
+    # Load Tiny-toxic-detector
+    config = TinyTransformerConfig.from_pretrained("AssistantsLab/Tiny-Toxic-Detector", use_auth_token=hf_token)
+    models["Tiny-toxic-detector"] = TinyTransformerForSequenceClassification.from_pretrained("AssistantsLab/Tiny-Toxic-Detector", config=config, use_auth_token=hf_token).to(device)
+    tokenizers["Tiny-toxic-detector"] = AutoTokenizer.from_pretrained("AssistantsLab/Tiny-Toxic-Detector", use_auth_token=hf_token)
+    
+    # Load other models
+    model_configs = [
+        ("unitary/toxic-bert", AutoModelForSequenceClassification, "unitary/toxic-bert"),
+        ("s-nlp/roberta_toxicity_classifier", AutoModelForSequenceClassification, "s-nlp/roberta_toxicity_classifier"),
+        ("martin-ha/toxic-comment-model", AutoModelForSequenceClassification, "martin-ha/toxic-comment-model"),
+        ("lmsys/toxicchat-t5-large-v1.0", AutoModelForSeq2SeqLM, "t5-large")
+    ]
+    
+    for model_name, model_class, tokenizer_name in model_configs:
+        models[model_name] = model_class.from_pretrained(model_name, use_auth_token=hf_token).to(device)
+        tokenizers[model_name] = AutoTokenizer.from_pretrained(tokenizer_name, use_auth_token=hf_token)
+    
+    return models, tokenizers, device
+
+# Prediction function
+def predict_toxicity(text, model, tokenizer, device, model_name):
+    start_time = time.time()
+    
+    if model_name == "lmsys/toxicchat-t5-large-v1.0":
+        prefix = "ToxicChat: "
+        inputs = tokenizer.encode(prefix + text, return_tensors="pt").to(device)
+        
+        with torch.no_grad():
+            outputs = model.generate(inputs, max_new_tokens=5)
+        
+        prediction = tokenizer.decode(outputs[0], skip_special_tokens=True).strip().lower()
+        prediction = "Toxic" if prediction == "positive" else "Not Toxic"
+    else:
+        inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128, padding="max_length").to(device)
+        
+        if "token_type_ids" in inputs:
+            del inputs["token_type_ids"]
+
+        with torch.no_grad():
+            outputs = model(**inputs)
+        
+        if model_name == "Tiny-toxic-detector":
+            logits = outputs["logits"].squeeze()
+            prediction = "Toxic" if logits > 0.5 else "Not Toxic"
+        else:
+            logits = outputs.logits.squeeze()
+            prediction = "Toxic" if logits[1] > logits[0] else "Not Toxic"
+
+    end_time = time.time()
+    inference_time = end_time - start_time
+    
+    return prediction, inference_time
+
+def main():
+    st.set_page_config(page_title="Multi-Model Toxicity Detector", layout="wide")
+    st.title("Multi-Model Toxicity Detector")
+
+    # Load models
+    hf_token = os.getenv('AT')
+    models, tokenizers, device = load_models_and_tokenizers(hf_token)
+
+    # Reorder the models dictionary so that "Tiny-toxic-detector" is last
+    model_names = sorted(models.keys(), key=lambda x: x == "Tiny-toxic-detector")
+
+    # User input
+    text = st.text_area("Enter text to classify:", height=150)
+
+    if st.button("Classify"):
+        if text:
+            progress_bar = st.progress(0)
+            results = []
+
+            for i, model_name in enumerate(model_names):
+                with st.spinner(f"Classifying with {model_name}..."):
+                    prediction, inference_time = predict_toxicity(text, models[model_name], tokenizers[model_name], device, model_name)
+                    results.append((model_name, prediction, inference_time))
+                progress_bar.progress((i + 1) / len(model_names))
+
+            st.success("Classification complete!")
+            progress_bar.empty()
+
+            # Display results in a grid
+            col1, col2, col3 = st.columns(3)
+            for i, (model_name, prediction, inference_time) in enumerate(results):
+                with [col1, col2, col3][i % 3]:
+                    st.subheader(model_name)
+                    st.write(f"Prediction: {prediction}")
+                    st.write(f"Inference Time: {inference_time:.4f}s")
+                    st.write("---")
+        else:
+            st.warning("Please enter some text to classify.")
+
+if __name__ == "__main__":
+    main()