from fastapi import FastAPI, Request
from transformers import pipeline
import os
import tempfile

app = FastAPI()

# Використання тимчасової директорії для кешу
cache_dir = tempfile.mkdtemp()
os.environ["TRANSFORMERS_CACHE"] = cache_dir
print(f"Використовується директорія кешу: {cache_dir}")

try:
    pipe = pipeline(
        "text-generation",
        model="google/gemma-3-1b-it",
        token=os.environ.get("HF_TOKEN"),  # Краще використовувати get для безпеки
        device="cpu",
        max_new_tokens=256,
        temperature=0.7
    )
except Exception as e:
    print(f"Помилка при ініціалізації pipeline: {e}")
    # Тут можна обробити помилку або підняти HTTP виняток

@app.post("/generate")
async def generate(request: Request):
    data = await request.json()
    prompt = data.get("prompt", "")
    try:
        result = pipe(prompt)
        return {"response": result[0]["generated_text"]}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))