Update app.py

app.py

@@ -1,22 +1,15 @@
 import gradio as gr
-from transformers import BertConfig, BertForSequenceClassification, AutoTokenizer
-from safetensors import safe_open
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import torch
+import pickle
 
-config_path = "peterkros/cofogv1-bert/modelbert2/config.json"
-config = BertConfig.from_json_file(config_path)
-model = BertForSequenceClassification(config)
-
-
-model_path = "modelbert2"
-model = BertForSequenceClassification.from_pretrained(model_path)
-tokenizer = AutoTokenizer.from_pretrained(model_path)
-
-tokenizer = AutoTokenizer.from_pretrained("peterkros/cofogv1-bert/modelbert2/")
+# Load the model and tokenizer from Hugging Face Hub
+model_name = "peterkros/cofogv1-bert/modelbert2"
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 # Load the label encoder
-import pickle
-with open('peterkros/cofogv1-bert/label_encoder.pkl', 'rb') as file:
+with open('label_encoder.pkl', 'rb') as file:
     label_encoder = pickle.load(file)
 
 def predict(text):
@@ -28,18 +21,17 @@ def predict(text):
     predicted_label = label_encoder.inverse_transform([predicted_class])[0]
     return predicted_label
 
-
 # Define the markdown text with bullet points
 markdown_text = """
-- Trainied with ~1500 rows of data on bert-base-uncased	110M, English
-- Input one budget line per time. 
+- Trained with ~1500 rows of data on bert-base-uncased, 110M, English.
+- Input one budget line per time.
 - Accuracy of the model is ~72%.
 """
 
 # Define the interface
 iface = gr.Interface(
     fn=predict, 
-    inputs=gr.inputs.Textbox(lines=1, placeholder="Enter Budget line here..."), 
+    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter Budget line here..."), 
     outputs="text", 
     title="COFOG Level 1 Classification",
     description=markdown_text  # Add the markdown text to the description
@@ -47,4 +39,4 @@ iface = gr.Interface(
 
 # Run the interface
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()