huimanho's picture
Update app.py
4b4ac3c verified
import streamlit as st
from transformers import AutoModelForSequenceClassification
from transformers import AutoTokenizer
from transformers import pipeline
import torch
import numpy as np
def main():
st.title("yelp2024fall Test")
st.write("Enter a sentence for analysis:")
user_input = st.text_input("")
if user_input:
# Approach: AutoModel
model2 = AutoModelForSequenceClassification.from_pretrained("huimanho/CustomModel_yelp",
num_labels=5)
sentiment_pipeline = pipeline(model="huimanho/CustomModel_yelp")
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
outputs = model2(**user_input)
predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
predictions = predictions.cpu().detach().numpy()
# Get the index of the largest output value
max_index = np.argmax(predictions)
st.write(f"result (AutoModel) - Label: {max_index}")
if __name__ == "__main__":
main()