Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -1,7 +1,6 @@
|
|
1 |
import streamlit as st
|
2 |
import os
|
3 |
from transformers import AutoModelForSequenceClassification, AutoTokenizer
|
4 |
-
#from transformers import GPT2Tokenizer, GPT2LMHeadModel
|
5 |
import torch
|
6 |
|
7 |
# Set CUDA_LAUNCH_BLOCKING for better debugging
|
@@ -22,12 +21,11 @@ label_mapping = {
|
|
22 |
"Campaign": 10,
|
23 |
"Intelligent Pigging": 11
|
24 |
}
|
|
|
25 |
model_name = "amiguel/names_scope_finetuned" # Replace with your model repo
|
26 |
model = AutoModelForSequenceClassification.from_pretrained(model_name)
|
27 |
-
#model = GPT2LMHeadModel.from_pretrained(model_name)
|
28 |
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
29 |
-
|
30 |
-
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
31 |
model.to(device)
|
32 |
|
33 |
def classify_text(text, model, tokenizer, device, max_length=512):
|
@@ -47,7 +45,10 @@ def classify_text(text, model, tokenizer, device, max_length=512):
|
|
47 |
logits = outputs.logits
|
48 |
predicted_label = torch.argmax(logits, dim=-1).item()
|
49 |
|
50 |
-
|
|
|
|
|
|
|
51 |
|
52 |
def main():
|
53 |
st.title("Scope vs ItemClass")
|
@@ -76,7 +77,7 @@ def main():
|
|
76 |
text_column = st.selectbox("Select the text column", df.columns)
|
77 |
predicted_labels = [classify_text(text, model, tokenizer, device) for text in df[text_column]]
|
78 |
df["Predicted Label"] = predicted_labels
|
79 |
-
st.write(df)
|
80 |
|
81 |
if __name__ == "__main__":
|
82 |
main()
|
|
|
1 |
import streamlit as st
|
2 |
import os
|
3 |
from transformers import AutoModelForSequenceClassification, AutoTokenizer
|
|
|
4 |
import torch
|
5 |
|
6 |
# Set CUDA_LAUNCH_BLOCKING for better debugging
|
|
|
21 |
"Campaign": 10,
|
22 |
"Intelligent Pigging": 11
|
23 |
}
|
24 |
+
|
25 |
model_name = "amiguel/names_scope_finetuned" # Replace with your model repo
|
26 |
model = AutoModelForSequenceClassification.from_pretrained(model_name)
|
|
|
27 |
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
28 |
+
device = torch.device("cpu") # Force using CPU for debugging
|
|
|
29 |
model.to(device)
|
30 |
|
31 |
def classify_text(text, model, tokenizer, device, max_length=512):
|
|
|
45 |
logits = outputs.logits
|
46 |
predicted_label = torch.argmax(logits, dim=-1).item()
|
47 |
|
48 |
+
if predicted_label < 0 or predicted_label >= len(label_mapping):
|
49 |
+
raise ValueError(f"Predicted label {predicted_label} is out of range.")
|
50 |
+
|
51 |
+
return list(label_mapping.keys())[predicted_label]
|
52 |
|
53 |
def main():
|
54 |
st.title("Scope vs ItemClass")
|
|
|
77 |
text_column = st.selectbox("Select the text column", df.columns)
|
78 |
predicted_labels = [classify_text(text, model, tokenizer, device) for text in df[text_column]]
|
79 |
df["Predicted Label"] = predicted_labels
|
80 |
+
st.write(df)
|
81 |
|
82 |
if __name__ == "__main__":
|
83 |
main()
|