Update app.py

app.py

@@ -3,15 +3,15 @@ from transformers import AutoTokenizer, FalconForSequenceClassification
 import torch
 
 # 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 = AutoTokenizer.from_pretrained("rocketknight1/falcon-rw-1b")
+model = FalconForSequenceClassification.from_pretrained("rocketknight1/falcon-rw-1b", num_labels=8)
 
 def predict(text):
     # Encode input text into IDs and mask
-    inputs = tokenizer(text, return_tensors='pt')
+    inputs = tokenizer(text, return_tensors="pt")
     
     # Run inference on encoded input
-    outputs = model(inputs['input_ids'], attention_mask=inputs['attention_mask'])
+    outputs = model(**inputs)
     
     # Get last hidden state (i.e., output of final layer)
     last_hidden_state = outputs.last_hidden_state[:, 0, :]
@@ -19,6 +19,6 @@ def predict(text):
     # Return predicted label
     return torch.argmax(last_hidden_state).item() + 1
 
-iface = gr.Interface(fn=predict, inputs="textbox", outputs="label", title="Falcon-7b")
+iface = gr.Interface(fn=predict, inputs="textbox", outputs="label", title="FALCON Chatbot")
 if __name__ == "__main__":
     iface.launch()