from flask import Flask, request, jsonify
from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForTokenClassification

# Initialize the tokenizer and model

app = Flask(__name__)
classifier = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base", return_all_scores=True)

@app.route('/classify', methods=['POST'])
def classify():
    try:
        data = request.get_json()
        if 'text' not in data:
            return jsonify({"error": "Missing 'text' field"}), 400
        
        text = data['text']
        result = classifier(text)
        return jsonify(result)
    
    except Exception as e:
        return jsonify({"error": str(e)}), 500

tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER")
model = AutoModelForTokenClassification.from_pretrained("dslim/bert-base-NER")
nlp = pipeline("ner", model=model, tokenizer=tokenizer)
@app.route('/ner', methods=['POST'])
def ner_endpoint():
    try:
        # Get text from request
        data = request.get_json()
        text = data.get("text", "")
        
        # Perform NER
        ner_results = nlp(text)
        
        # Extract words and their corresponding entities
        words_and_entities = [
            {"word": result['word'], "entity": result['entity']} 
            for result in ner_results
        ]
        
        # Return JSON response with the words and their entities
        return jsonify({"entities": words_and_entities})
    except Exception as e:
        return jsonify({"error": str(e)}), 500