huimanho's picture
Update app.py
bfbb65f verified
raw
history blame
1.05 kB
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(**inputs)
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()