from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
import torch

app = FastAPI()

# Model configuration
MODEL_NAME = "nlptown/bert-base-multilingual-uncased-sentiment"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Initialize model and tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
classifier = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer, device=DEVICE)

class TextInput(BaseModel):
    text: str

@app.post("/analyze-sentiment")
async def analyze_sentiment(input_data: TextInput):
    try:
        result = classifier(input_data.text)
        return {
            "sentiment": result[0]['label'],
            "score": float(result[0]['score'])
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Przykład dla większego modelu (np. GPT-2)
MODEL_NAME_LARGE = "gpt2-large"
tokenizer_large = AutoTokenizer.from_pretrained(MODEL_NAME_LARGE)
model_large = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME_LARGE)

class GenerationInput(BaseModel):
    prompt: str
    max_length: int = 100

@app.post("/generate-text")
async def generate_text(input_data: GenerationInput):
    try:
        inputs = tokenizer_large(input_data.prompt, return_tensors="pt")
        outputs = model_large.generate(
            inputs["input_ids"],
            max_length=input_data.max_length,
            num_return_sequences=1,
            no_repeat_ngram_size=2
        )
        generated_text = tokenizer_large.decode(outputs[0], skip_special_tokens=True)
        return {"generated_text": generated_text}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Dodanie podstawowego health checka
@app.get("/health")
async def health_check():
    return {"status": "healthy"}