Update app.py

app.py

@@ -1,22 +1,24 @@
 import gradio as gr
-from transformers import AutoTokenizer, FalconModel
+from transformers import AutoTokenizer, FalconForSequenceClassification
+import torch
 
-tokenizer = AutoTokenizer.from_pretrained("Rocketknight1/falcon-rw-1b")
-model = FalconModel.from_pretrained("Rocketknight1/falcon-rw-1b")
+# 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)
 
-def chat(input_text):
-    inputs = tokenizer(input_text, return_tensors="pt")
-    outputs = model(**inputs)
-    last_hidden_states = outputs.last_hidden_state
-    response = tokenizer.decode(last_hidden_states.squeeze(), skip_special_tokens=True)
-    return response
+def predict(text):
+    # Encode input text into IDs and mask
+    inputs = tokenizer(text, return_tensors='pt')
+    
+    # Run inference on encoded input
+    outputs = model(inputs['input_ids'], attention_mask=inputs['attention_mask'])
+    
+    # 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
 
-interface = gr.Interface(
-    fn=chat,
-    inputs=gr.Textbox(label="User Question"),
-    outputs=gr.Textbox(label="Chatbot Response"),
-    title="Falcon-7b",
-    theme="dark"
-)
-
-interface.launch()
+iface = gr.Interface(fn=predict, inputs="textbox", outputs="label", title="Falcon-7b")
+if __name__ == "__main__":
+    iface.launch()