test / app.py
TInnu's picture
Create app.py
2cbdb7c verified
raw
history blame contribute delete
907 Bytes
import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification
# Load TinyBERT
model_name = "huawei-noah/TinyBERT_General_6L_768D"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
# Streamlit App Title
st.title("TinyBERT Text Summarization")
# Input text box
input_text = st.text_area("Enter text for summarization:", height=200)
# Button to perform summarization
if st.button("Summarize"):
if input_text:
# Tokenize input text
inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
# Get model outputs
outputs = model(**inputs)
# Display output (this is placeholder logic, adjust to your specific task)
st.write(f"Model output: {outputs}")
else:
st.warning("Please enter some text to summarize.")