import os import uvicorn import torch from fastapi import FastAPI, HTTPException from fastapi.responses import JSONResponse, Response from pydantic import BaseModel from transformers import pipeline # Utwórz instancję FastAPI app = FastAPI( title="Bielik Text Generation API", description="API do generowania tekstu za pomocą modelu Bielik-1.5B-v3.0-Instruct", version="1.0.0" ) # Ścieżka do modelu - Hugging Face automatycznie pobierze model MODEL_NAME = "speakleash/Bielik-1.5B-v3.0-Instruct" generator = None # Zostanie załadowany później # Model wejściowy dla POST request class GenerationRequest(BaseModel): prompt: str max_new_tokens: int = 50 temperature: float = 0.7 top_p: float = 0.9 @app.on_event("startup") async def startup_event(): """ Ładowanie modelu podczas uruchamiania aplikacji. To zajmie trochę czasu, ale dzieje się tylko raz. """ global generator print(f"Ładowanie modelu: {MODEL_NAME}...") try: # Możesz dostosować device=0 (GPU) lub device=-1 (CPU) w zależności od wybranej maszyny Space # Free tier spaces usually run on CPU, unless you explicitly select a GPU. # It's safer to not specify device if you want it to auto-detect or default to CPU. generator = pipeline( "text-generation", model=MODEL_NAME, # device=0 if torch.cuda.is_available() else -1 # Odkomentuj dla detekcji GPU ) print("Model załadowany pomyślnie!") except Exception as e: print(f"Błąd ładowania modelu: {e}") # Możesz zdecydować, czy aplikacja ma zakończyć działanie, czy kontynuować bez modelu # W przypadku błędu ładowania modelu, endpoint generacji tekstu będzie zwracał błąd generator = None # Ustaw na None, aby sygnalizować problem @app.get("/") async def root(): """ Główny endpoint (health check). """ return {"message": "Bielik Text Generation API is running!"} @app.post("/generate") async def generate_text(request: GenerationRequest): """ Endpoint do generowania tekstu na podstawie promptu. """ if generator is None: raise HTTPException(status_code=503, detail="Model nie został załadowany lub wystąpił błąd.") try: generated_text = generator( request.prompt, max_new_tokens=request.max_new_tokens, temperature=request.temperature, top_p=request.top_p, do_sample=True, # Ważne dla generowania z temperaturą return_full_text=False # Dodaj ten parametr, aby model nie zwracał od razu promptu ) # Pipeline zwraca listę słowników, bierzemy pierwszy wynik #gen_text = {"generated_text": generated_text[0]["generated_text"]} response_data = {"generated_text": generated_text[0]["generated_text"]} return JSONResponse( content=response_data, media_type="application/json; charset=utf-8" ) #return Response(content=generated_text[0]["generated_text"], media_type="text/plain; charset=utf-8") # return {"generated_text": generated_text[0]["generated_text"]} except Exception as e: raise HTTPException(status_code=500, detail=f"Błąd podczas generowania tekstu: {e}") # Uruchamianie serwera Uvicorn bezpośrednio (dla Dockera) if __name__ == "__main__": uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)))