Update app.py

app.py

@@ -1,24 +1,15 @@
 import gradio as gr
-from transformers import AutoTokenizer, FalconForSequenceClassification
-import torch
+from transformers import BertTokenizer, BertForSequenceClassification
 
-# Load pre-trained FALCON model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("rocketknight1/falcon-rw-1b")
-model = FalconForSequenceClassification.from_pretrained("rocketknight1/falcon-rw-1b", num_labels=8)
+tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=8)
 
-def predict(text):
-    # Encode input text into IDs and mask
-    inputs = tokenizer(text, return_tensors="pt")
-    
-    # Run inference on encoded input
-    outputs = model(**inputs)
-    
-    # Get last hidden state (i.e., output of final layer)
-    last_hidden_state = outputs.last_hidden_state[:, 0, :]
-    
-    # Return predicted label
-    return torch.argmax(last_hidden_state).item() + 1
+def classify(text):
+    inputs = tokenizer(text, return_tensors='pt')
+    outputs = model(inputs['input_ids'], attention_mask=inputs['attention_mask'])
+    logits = outputs[0].logits
+    probabilities = torch.softmax(logits, dim=-1).tolist()
+    predicted_label = np.argmax(probabilities)
+    return {'Label': predicted_label}
 
-iface = gr.Interface(fn=predict, inputs="textbox", outputs="label", title="FALCON Chatbot")
-if __name__ == "__main__":
-    iface.launch()
+gr.Interface(fn=classify, inputs="text", outputs="json").launch()