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

@app.post("/emotion")
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

@app.post("/ner")
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})