Update app.py

app.py

@@ -1,19 +1,28 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, BartModel
 import torch
 
-# Load the BART tokenizer and model
+# Load the tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
-model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-base")
+model = BartModel.from_pretrained("facebook/bart-base")
 
-def generate_response(user_input):
-    inputs = tokenizer(user_input, return_tensors="pt", max_length=512, truncation=True)
-    outputs = model.generate(**inputs)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+def generate_response(text):
+    # Tokenize the input text
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
 
-input_textbox = gr.Textbox(lines=10, label="Enter your text here")
-output_textbox = gr.Textbox(label="Chatbot Response")
+    # Generate model output
+    outputs = model(**inputs)
 
-chatbot_interface = gr.Interface(fn=generate_response, inputs=input_textbox, outputs=output_textbox, title="Hugging Face BART Chatbot", description="This chatbot uses the Hugging Face BART model to generate responses based on user input.")
-chatbot_interface.launch()
+    # Get the last hidden states from the model output
+    last_hidden_states = outputs.last_hidden_state
+
+    # Return last hidden states as output
+    return last_hidden_states
+
+iface = gr.Interface(
+    fn=generate_response,
+    inputs=gr.Textbox(lines=7, label="Input Text"),
+    outputs=gr.Textbox(label="Output Last Hidden States")
+)
+
+iface.launch()