Spaces:
Sleeping
Sleeping
| from fastapi import FastAPI | |
| from pydantic import BaseModel | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
| import torch | |
| # Define the model and tokenizer | |
| tokenizer_emotion = AutoTokenizer.from_pretrained("SamLowe/roberta-base-go_emotions") | |
| model_emotion = AutoModelForSequenceClassification.from_pretrained("SamLowe/roberta-base-go_emotions") | |
| from transformers import AutoModelForTokenClassification | |
| tokenizer_t_class = AutoTokenizer.from_pretrained("dslim/bert-base-NER") | |
| model_t_class = AutoModelForTokenClassification.from_pretrained("dslim/bert-base-NER") | |
| # FastAPI app | |
| app = FastAPI() | |
| # Define the request body | |
| class Input(BaseModel): | |
| text: str | |
| async def predict_emotion(input: Input): | |
| # Tokenize the input text | |
| inputs = tokenizer_emotion(input.text, return_tensors="pt") | |
| # Get the model's predictions | |
| outputs = model_emotion(**inputs) | |
| # Get the predicted class | |
| predicted_class_idx = torch.argmax(outputs.logits).item() | |
| # Decode the predicted class | |
| predicted_class = model_emotion.config.id2label[predicted_class_idx] | |
| # Return the prediction | |
| return {"predicted_emotion": predicted_class} | |
| # from transformers import pipeline | |
| async def perform_ner(item: Input): | |
| # nlp = pipeline("ner", model=model_t_class, tokenizer=tokenizer_t_class) | |
| # ner_results = nlp(item.text) | |
| # print(ner_results) | |
| # Tokenize the input text | |
| inputs = tokenizer_t_class.encode(item.text, return_tensors="pt") | |
| # Get the model's predictions | |
| outputs = model_t_class(inputs).logits | |
| # Get the predicted classes for each token | |
| predictions = torch.argmax(outputs, dim=2) | |
| # Decode the tokens and their predicted classes | |
| tokens = tokenizer_t_class.convert_ids_to_tokens(inputs[0]) | |
| predicted_labels = [model_t_class.config.id2label[prediction] for prediction in predictions[0].tolist()] | |
| # Pair tokens with their predicted labels | |
| token_labels = list(zip(tokens, predicted_labels)) | |
| # Filter out subwords | |
| filtered_token_labels = [(token, label) for token, label in token_labels if not token.startswith("##")] | |
| # Return the tokens and their predicted labels | |
| return {"token_labels": filtered_token_labels} | |
| # if __name__ == "__main__": | |
| # import uvicorn | |
| # uvicorn.run(app, host="0.0.0.0", port=8000) | |
| import streamlit as st | |
| st.title("Emotion Detection App") | |
| input_text = st.text_input("Enter your text here") | |
| if st.button("Predict"): | |
| inputs = tokenizer_emotion(input_text, return_tensors="pt") | |
| # Get the model's predictions | |
| outputs = model_emotion(**inputs) | |
| # Get the predicted class | |
| predicted_class_idx = torch.argmax(outputs.logits).item() | |
| # Decode the predicted class | |
| predicted_class = model_emotion.config.id2label[predicted_class_idx] | |
| # Return the prediction | |
| st.write({"predicted_emotion": predicted_class}) | |
| # Tokenize the input text | |
| inputs = tokenizer_t_class.encode(input_text, return_tensors="pt") | |
| # Get the model's predictions | |
| outputs = model_t_class(inputs).logits | |
| # Get the predicted classes for each token | |
| predictions = torch.argmax(outputs, dim=2) | |
| # Decode the tokens and their predicted classes | |
| tokens = tokenizer_t_class.convert_ids_to_tokens(inputs[0]) | |
| predicted_labels = [model_t_class.config.id2label[prediction] for prediction in predictions[0].tolist()] | |
| # Pair tokens with their predicted labels | |
| token_labels = list(zip(tokens, predicted_labels)) | |
| # Filter out subwords | |
| filtered_token_labels = [(token, label) for token, label in token_labels if not token.startswith("##")] | |
| # Return the tokens and their predicted labels | |
| st.write({"token_labels": filtered_token_labels}) |