more fixes

app.py

@@ -15,10 +15,15 @@ model = BertForSequenceClassification.from_pretrained("sleiyer/restricted_item_d
 # Load the trained model and tokenizer
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
 
+from pydantic import BaseModel
+
+class Predict(BaseModel):
+  input: str
+
 # Function to predict the class of a single input text
-def predict(text: str):
+def predict(request: Predict):
   # Preprocess the input text
-  inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True)
+  inputs = tokenizer(request.input, return_tensors='pt', truncation=True, padding=True)
 
   # Make predictions
   with torch.no_grad():
@@ -34,10 +39,10 @@ def predict(text: str):
   predicted_label = label_map[predicted_class]
 
   # Displaying the user input
-  return f'The item "{text}" is classified as: "{predicted_label}"'
+  return f'The item "{request.input}" is classified as: "{predicted_label}"'
 
   return predicted_class
 
 @app.post("/predict")
-def predict(input):
-  return predict(input)
+def predictApi(request: Predict):
+  return predict(request)

restrictedItems/predict.py

@@ -2,7 +2,8 @@ from transformers import BertTokenizer, BertForSequenceClassification
 import torch
 
 # Load the trained model and tokenizer
-model = BertForSequenceClassification.from_pretrained("/Users/slei/hackweek2024-sup-genai-tools/hackweek2024-sup-genai-tools/trained_model")
+# 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

restrictedItems/train.py

@@ -100,7 +100,7 @@ train_dataset = ShoppingCartDataset(train_encodings, train_labels)
 val_dataset = ShoppingCartDataset(val_encodings, val_labels)
 
 # Load pre-trained BERT model
-model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
+model = BertForSequenceClassification.from_pretrained('sleiyer/restricted_item_detector')
 
 # Training arguments
 training_args = TrainingArguments(
@@ -128,4 +128,5 @@ trainer.train()
 # Evaluate model
 trainer.evaluate()
 
+model.save_pretrained('trained_model')
 model.push_to_hub("sleiyer/restricted_item_detector")