from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
import torch

# Inisialisasi model dan tokenizer
model_name = "ragilbuaj/sentiment-analysis-TWS-reviews"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Inisialisasi FastAPI
app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Bisa disesuaikan dengan daftar asal yang diizinkan
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Model request body
class TextInput(BaseModel):
    text: str

# Fungsi untuk analisis sentimen
def predict_sentiment(text):
    nlp = pipeline(
        "sentiment-analysis",
        model=model_name,
        tokenizer=model_name
    )

    result = nlp(text)[0]
    sentiment = result['label']
    confidence = result['score']
    return sentiment, confidence



# Endpoint untuk analisis sentimen
@app.post("/predict")
async def predict(input: TextInput):
    sentiment, confidence = predict_sentiment(input.text)
    return {"sentiment": sentiment, "confidence": confidence}

# Endpoint root
@app.get("/")
async def read_root():
    return {"message": "Sentiment Analysis API"}