Update app.py

app.py

@@ -1,67 +1,49 @@
-import gradio as gr
 import torch
-from transformers import BertForSequenceClassification, BertTokenizer
+from transformers import AutoModel, AutoTokenizer
+import torch.nn.functional as F
 
-# Load the tokenizer from Hugging Face
-token_model = "indolem/indobertweet-base-uncased"
-tokenizer = BertTokenizer.from_pretrained(token_model)
+def load_model(model_name):
+    # Load model from Hugging Face Hub
+    model = AutoModel.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    return model, tokenizer
 
-# Define the model directory where your config.json and pytorch_model.bin are located
-model_directory = "model_directory"  # Make sure this directory has config.json and pytorch_model.bin
-
-# Load the model
-# If your weights are named differently, ensure the file is named pytorch_model.bin or modify the loading method
-model = BertForSequenceClassification.from_pretrained(model_directory)
-model.eval()  # Set the model to evaluation mode
-
-# Check if CUDA is available and set the device accordingly
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model.to(device)
-
-def classify_transaction(notes):
-    # Tokenize the input text
-    inputs = tokenizer.encode_plus(
-        notes,
-        None,
-        add_special_tokens=True,
-        max_length=256,
-        padding='max_length',
-        return_token_type_ids=False,
-        return_attention_mask=True,
-        truncation=True,
-        return_tensors='pt'
-    )
+def predict(model, tokenizer, input_text):
+    # Preprocess the input
+    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
 
     # Move tensors to the same device as the model
-    input_ids = inputs['input_ids'].to(device)
-    attention_mask = inputs['attention_mask'].to(device)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    inputs = {k: v.to(device) for k, v in inputs.items()}
 
-    # Model in evaluation mode
+    # Make the model generate a prediction
     model.eval()
-
-    # Make prediction
     with torch.no_grad():
-        outputs = model(input_ids, attention_mask=attention_mask)
+        outputs = model(**inputs)
+        logits = outputs.logits
 
-    # Extract logits and convert to probabilities
-    logits = outputs[0]
-    probabilities = torch.softmax(logits, dim=1)
+    # Convert logits to probabilities
+    probabilities = F.softmax(logits, dim=1)
 
     # Get the predicted class
     predicted_class = torch.argmax(probabilities, dim=1).cpu().numpy()
+    probabilities = probabilities.cpu().numpy()
+
+    return predicted_class, probabilities
 
-    # Return the predicted class
-    return f"Predicted Category: {predicted_class}"
+def main():
+    model_name = "your-huggingface-model-name"  # Replace with your model's name
+    model, tokenizer = load_model(model_name)
 
-# Creating the Gradio interface
-iface = gr.Interface(
-    fn=classify_transaction,
-    inputs=gr.Textbox(lines=3, placeholder="Enter Transaction Notes Here", label="Transaction Notes"),
-    outputs=gr.Text(label="Classification Result"),
-    title="Transaction Category Classifier",
-    description="Enter transaction notes to get the predicted category.",
-    live=True  # Update the output as soon as the input changes
-)
+    # Example input
+    input_text = "Your input text goes here"
+    
+    # Get predictions
+    predicted_class, probabilities = predict(model, tokenizer, input_text)
+    
+    print("Predicted Class:", predicted_class)
+    print("Probabilities:", probabilities)
 
 if __name__ == "__main__":
-    iface.launch()
+    main()