new app

app.py

@@ -1,77 +1,23 @@
 import streamlit as st
 import json
 import numpy as np
+import sentencepiece
+from pathlib import Path
 import torch
 from transformers import (
     DebertaV2Config,
     DebertaV2Model,
     DebertaV2Tokenizer,
 )
-import sentencepiece 
 
-try:
-    import sentencepiece
-except ImportError:
-    st.error("Требуется установить SentencePiece: pip install sentencepiece")
-    st.stop()
+MODEL_NAME = "microsoft/deberta-v3-base"
+MAX_LENGTH = 512
+NUM_LABELS = 47
+DROPOUT_RATE = 0.1
+THRESHOLD = 0.5
 
-model_name = "microsoft/deberta-v3-base"
-tokenizer = DebertaV2Tokenizer.from_pretrained(model_name)
-
-def preprocess_text(text, tokenizer, max_length=512):
-    inputs = tokenizer(
-        text,
-        padding="max_length",
-        truncation=True,
-        max_length=max_length,
-        return_tensors="pt"
-    )
-    return inputs
-
-
-def classify_text(text, model, tokenizer, device, threshold=0.5):
-    inputs = preprocess_text(text, tokenizer)
-    input_ids = inputs["input_ids"].to(device)
-    attention_mask = inputs["attention_mask"].to(device)
-    model.eval()
-    with torch.no_grad():
-        logits = model(input_ids, attention_mask)
-    probs = torch.sigmoid(logits)
-    predictions = (probs > threshold).int().cpu().numpy()
-    
-    return probs.cpu().numpy(), predictions
-
-def get_themes(text, model, tokenizer, label_to_theme, device, limit=5):
-    probabilities, _ = classify_text(text, model, tokenizer, device)
-    print(probabilities)
-    themes = []
-    for label in probabilities[0].argsort()[-limit:]:
-        themes.append((label_to_theme[str(label)], probabilities[0][label]))
-    return themes
-
-class DebertPaperClassifier(torch.nn.Module):
-    def __init__(self, num_labels, device, dropout_rate=0.1, class_weights=None):
-        super().__init__()
-        self.config = DebertaV2Config.from_pretrained(model_name)
-        self.deberta = DebertaV2Model.from_pretrained(model_name, config=self.config)
-
-        self.classifier = torch.nn.Sequential(
-            torch.nn.Dropout(dropout_rate),
-            torch.nn.Linear(self.config.hidden_size, 512),
-            torch.nn.LayerNorm(512),
-            torch.nn.GELU(),
-            torch.nn.Dropout(dropout_rate),
-            torch.nn.Linear(512, num_labels)
-        )
-
-        self._init_weights()
-        if class_weights is not None:
-            self.loss_fct = torch.nn.BCEWithLogitsLoss(weight=class_weights.to(device))
-        else:
-            self.loss_fct = torch.nn.BCEWithLogitsLoss()
-
-class DebertPaperClassifierV5(torch.nn.Module):
-    def __init__(self, device, num_labels=47, dropout_rate=0.1, class_weights=None):
+class DebertaV3PaperClassifier(torch.nn.Module):
+    def __init__(self, device, num_labels=NUM_LABELS, dropout_rate=DROPOUT_RATE, class_weights=None):
         super().__init__()
         self.config = DebertaV2Config.from_pretrained("microsoft/deberta-v3-base")
         self.deberta = DebertaV2Model.from_pretrained("microsoft/deberta-v3-base", config=self.config)
@@ -107,48 +53,88 @@ class DebertPaperClassifierV5(torch.nn.Module):
                 module.weight.data.normal_(mean=0.0, std=0.02)
                 if module.bias is not None:
                     module.bias.data.zero_()
-    
+
 @st.cache_resource
-def load_model(test=False):
+def load_assets():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    try:
-        path = '/Users/bvd757/my_documents/Машинное обучение 2/Howework4'
-        with open(f'{path}/label_to_theme.json', 'r') as f:
-            label_to_theme = json.load(f)
-    except:
-        path = '/home/user/app'
-        with open(f'{path}/label_to_theme.json', 'r') as f:
-            label_to_theme = json.load(f)
-
-    # model = DebertPaperClassifier(device=device, num_labels=len(label_to_theme), class_weights=class_weights).to(device)
-    # model.load_state_dict(torch.load("model_info/full_model_v4.pth", map_location=device))
-
-    class_weights = torch.load(f'{path}/class_weights.pth').to(device)
-    model = DebertPaperClassifierV5(device=device, num_labels=47, class_weights=class_weights).to(device)
-    model.load_state_dict(torch.load(f"{path}/full_model_v4.pth", map_location=device))
-    if test:
-        print(device)
-        print("Model!!!")
-        text = 'We propose an architecture for VQA which utilizes recurrent layers to\ngenerate visual and textual attention. The memory characteristic of the\nproposed recurrent attention units offers a rich joint embedding of visual and\ntextual features and enables the model to reason relations between several\nparts of the image and question. Our single model outperforms the first place\nwinner on the VQA 1.0 dataset, performs within margin to the current\nstate-of-the-art ensemble model. We also experiment with replacing attention\nmechanisms in other state-of-the-art models with our implementation and show\nincreased accuracy. In both cases, our recurrent attention mechanism improves\nperformance in tasks requiring sequential or relational reasoning on the VQA\ndataset.'
-        print(get_themes(text, model, tokenizer, label_to_theme, device))
+
+    base_path = Path(__file__).parent
+    with open(base_path/"label_to_theme.json") as f:
+        label_to_theme = json.load(f)
+
+    class_weights = torch.load(f"{base_path}/class_weights.pth").to(device)
+
+    model = DebertaV3PaperClassifier(device=device, num_labels=NUM_LABELS, class_weights=class_weights).to(device)
+    model.load_state_dict(torch.load(f"{base_path}/full_model_v4.pth", map_location=device))
+    model.eval()
+
+    tokenizer = DebertaV2Tokenizer.from_pretrained(MODEL_NAME)
     return model, tokenizer, label_to_theme, device
 
 
-def kek():
+def preprocess_text(text, tokenizer, max_length=MAX_LENGTH):
+    inputs = tokenizer(
+        text,
+        padding="max_length",
+        truncation=True,
+        max_length=max_length,
+        return_tensors="pt"
+    )
+    return inputs
+
+
+def predict(text: str, model, tokenizer, device) -> list:
+    """Run model prediction on input text."""
+    inputs = preprocess_text(text, tokenizer)
+    
+    with torch.no_grad():
+        logits = model(
+            input_ids=inputs["input_ids"].to(device),
+            attention_mask=inputs["attention_mask"].to(device)
+        )
+    
+    probs = torch.sigmoid(logits).cpu().numpy()[0]
+    return probs
+
+
+def get_themes(probs: np.ndarray, label_to_theme: dict) -> list:
+    """Get top K themes with probabilities."""
+    sorted_indices = np.argsort(-probs)
+    labels = []
+    sum_percent = 0
+    for idx in sorted_indices:
+        labels.append((label_to_theme[str(idx)], probs[idx]))
+        sum_percent += probs[idx]
+        if sum_percent >= 0.95:
+            break
     
+    return labels
+
+
+def main():
+    st.title("Paper Classification App")
+    st.write("Classify research papers using DeBERTa model")
+
+    model, tokenizer, label_to_theme, device = load_assets()
+
     title = st.text_input("Title")
     abstract = st.text_area("Abstract")
-    
+
     if st.button("Classify"):
         if not title and not abstract:
-            st.warning("Please enter an abstract")
+            st.warning("Please enter title and/or abstract")
             return
         
-        text = f"{title}\n\n{abstract}" if title and abstract else title or abstract
-        model, tokenizer, label_to_theme, device = load_model()
-        
-        with st.spinner("Classifying..."):
-            themes = get_themes(text, model, tokenizer, label_to_theme, device, 5)
+        if abstract is None:
+            text = title
+        elif text is None:
+            text = abstract
+        else:
+            text = f"{title}\n\n{abstract}"
+
+        with st.spinner("Analyzing text..."):
+            probabilities = predict(text, model, tokenizer, device)
+            themes = get_themes(probabilities, label_to_theme)
         
         st.success("Classification results:")
         for theme, prob in themes:
@@ -156,4 +142,4 @@ def kek():
 
 
 if __name__ == "__main__":
-    kek()
+    main()