Add application file

app.py

@@ -0,0 +1,83 @@
+import gradio as gr
+import torch
+import re
+import nltk
+from nltk.stem import WordNetLemmatizer
+from nltk.corpus import stopwords
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+
+# Modeļu inicializācija
+model_names = ["distilbert-base-uncased", "prajjwal1/bert-tiny", "roberta-base", "google/mobilebert-uncased", "albert-base-v2", "xlm-roberta-base"]
+
+models = {}
+tokenizers = {}
+
+for model_name in model_names:
+    # Tokenizators
+    tokenizers[model_name] = AutoTokenizer.from_pretrained(model_name, max_length=512)
+
+    # Modelis ar 3 klasēm
+    models[model_name] = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=3)
+
+    model_file_name = re.sub(r'/', '_', model_name)
+    models[model_name].load_state_dict(torch.load(f"best_model_{model_file_name}.pth", map_location=torch.device('cpu')))
+    
+    # Uz ierīces
+    models[model_name] = models[model_name].to('cpu')
+    models[model_name].eval()
+
+# Label mapping
+labels = {0: "Safe", 1: "Spam", 2: "Phishing"}
+
+lemmatizer = WordNetLemmatizer()
+stop_words = set(stopwords.words('english'))
+
+def preprocess(text):
+    text = text.lower()  # Teksta pārveide atmetot lielos burtus
+    text = re.sub(r'http\S+', '', text)  # URL atmešana
+    text = re.sub(r"[^a-z']", ' ', text)  # atmet simbolus, kas nav burti
+    text = re.sub(r'\s+', ' ', text).strip()  # atmet liekās atstarpes
+    text = ' '.join([lemmatizer.lemmatize(word) for word in text.split() if word not in stop_words])  # lemmatizācija
+    return text
+
+# Classification function (single model)
+def classify_email_single_model(text, model_name):
+    text = preprocess(text)
+    inputs = tokenizers[model_name](text, return_tensors="pt", padding=True, truncation=True)
+    with torch.no_grad():
+        outputs = models[model_name](**inputs)
+        prediction = torch.argmax(outputs.logits, dim=1).item()
+    return labels[prediction]
+
+# Classification function (all models together)
+def classify_email(text):
+    votes = {"Safe": 0, "Spam": 0, "Phishing": 0}
+    
+    for model_name in model_names:
+        vote = classify_email_single_model(text, model_name)
+        votes[vote] += 1
+        
+        
+    response = ""
+    i = 1
+    for label, vote_count in votes.items():
+        vote_or_votes = "vote" if vote_count == 1 else "votes"
+        if i != 3:
+            response += f"{label}: {vote_count} {vote_or_votes}, "
+        else:
+            response += f"{label}: {vote_count} {vote_or_votes}"
+        i += 1
+        
+    return response
+    
+# Gradio UI
+demo = gr.Interface(
+    fn=classify_email,
+    inputs=gr.Textbox(lines=10, placeholder="Ievietojiet savu e-pastu šeit..."),
+    outputs="text",
+    title="E-pastu klasifikators (vairāku modeļu balsošana)",
+    description="Autori: Kristaps Tretjuks un Aleksejs Gorlovičs"
+)
+
+demo.launch(share=True)