Update script

script

@@ -1,26 +1,42 @@
 import torch
-from transformers import DistilBertTokenizer, DistilBertModel
+from transformers import DistilBertModel, DistilBertTokenizer
 
 # Load the tokenizer and model
-tokenizer = DistilBertTokenizer.from_pretrained("tokenizer_config.json")
-model = DistilBertModel.from_pretrained("model.pt")
+tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
+model = DistilBertCNN(num_labels=3)  # Assuming you have defined the custom classification layers
 
-# Define the inference function
-def predict(text):
-    # Tokenize the input
-    inputs = tokenizer(text, padding="max_length", truncation=True, return_tensors="pt")
+# Move the model to CPU
+device = torch.device("cpu")
+model.to(device)
 
-    # Perform the inference
-    with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
+# Load the saved model state dictionary
+model.load_state_dict(torch.load("path/to/save/directory/model.pt", map_location=device))
 
-    # Convert logits to probabilities
-    probabilities = torch.softmax(logits, dim=1).squeeze().tolist()
+# Set the model to evaluation mode
+model.eval()
 
-    return probabilities
+# Define a function to predict the class of a given tweet
+def classify_tweet(tweet):
+    inputs = tokenizer.encode_plus(
+        tweet,
+        add_special_tokens=True,
+        max_length=128,
+        padding="max_length",
+        truncation=True,
+        return_tensors="pt"
+    )
+    input_ids = inputs["input_ids"].to(device)
+    attention_mask = inputs["attention_mask"].to(device)
+    
+    with torch.no_grad():
+        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+    
+    logits = outputs[0]
+    predicted_class = torch.argmax(logits).item()
+    
+    return predicted_class
 
 # Example usage
-text = "This is a sample input."
-probabilities = predict(text)
-print(probabilities)
+tweet = "This is a sample tweet."
+predicted_class = classify_tweet(tweet)
+print(f"Predicted Class: {predicted_class}")