Spaces:
Sleeping
Sleeping
File size: 1,274 Bytes
28de1fd 2a6f3b3 28de1fd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
from transformers import BertTokenizer, BertForSequenceClassification
import torch
# Load the trained model and tokenizer
# model = BertForSequenceClassification.from_pretrained("/Users/slei/hackweek2024-sup-genai-tools/spaces/restricted_item_detector/trained_model")
model = BertForSequenceClassification.from_pretrained("sleiyer/restricted_item_detector")
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
# Function to predict the class of a single input text
def predict(text):
# Preprocess the input text
inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True)
# Make predictions
with torch.no_grad():
outputs = model(**inputs)
# Get the predicted class
logits = outputs.logits
predicted_class = torch.argmax(logits, dim=1).item()
return predicted_class
label_map = {0: 'Allowed Item', 1: 'Restricted Item'}
def main():
while True:
# Prompting the user for input
user_input = input("Enter something: ")
predicted_class = predict(user_input)
# Map the predicted class to a human-readable label
predicted_label = label_map[predicted_class]
# Displaying the user input
print(f'The item "{user_input}" is classified as: "{predicted_label}"')
if __name__ == "__main__":
main()
|