get ready for branching off for hf

language_models_project/app.py

@@ -22,6 +22,7 @@ model_name = st.selectbox(
         "finiteautomata/bertweet-base-sentiment-analysis",
         "ahmedrachid/FinancialBERT-Sentiment-Analysis",
         "finiteautomata/beto-sentiment-analysis",
+        "NativeVex/custom-fine-tuned"
     ],
 )
 
@@ -29,6 +30,38 @@ input_sentences = st.text_area("Sentences", value=demo_phrases, height=200)
 
 data = input_sentences.split("\n")
 
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+model_path = "bin/model4"
+model = AutoModelForSequenceClassification.from_pretrained(model_path)
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+from typing import List
+import torch
+import numpy as np
+import pandas as pd
+
+def infer(text: str) -> List[float]:
+    encoding = tokenizer(text, return_tensors="pt")
+    encoding = {k: v.to(model.device) for k,v in encoding.items()}
+    outputs = model(**encoding)
+    logits = outputs.logits
+    sigmoid = torch.nn.Sigmoid()
+    probs = sigmoid(logits.squeeze().cpu())
+    predictions = np.zeros(probs.shape)
+    predictions[np.where(probs >= 0.5)] = 1
+    predictions = pd.Series(predictions == 1)
+    predictions.index = ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"]
+    return predictions
+
+
+def wrapper(*args, **kwargs):
+    if args[0] != "NativeVex/custom-fine-tuned":
+        return classify(*args, **kwargs)
+    else:
+        return infer(text=args[1])
+
+    
+
 if st.button("Classify"):
     st.write("Please allow a few minutes for the model to run/download")
     for i in range(len(data)):