Upload main.py

main.py

@@ -1,8 +1,10 @@
+import os
+
+import uvicorn
 from fastapi import FastAPI, HTTPException
+from fastapi.responses import Response
 from pydantic import BaseModel
 from transformers import pipeline
-import uvicorn
-import os
 
 # Utwórz instancję FastAPI
 app = FastAPI(
@@ -13,7 +15,8 @@ app = FastAPI(
 
 # Ś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
+generator = None  # Zostanie załadowany później
+
 
 # Model wejściowy dla POST request
 class GenerationRequest(BaseModel):
@@ -22,6 +25,7 @@ class GenerationRequest(BaseModel):
     temperature: float = 0.7
     top_p: float = 0.9
 
+
 @app.on_event("startup")
 async def startup_event():
     """
@@ -44,7 +48,7 @@ async def startup_event():
         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
+        generator = None  # Ustaw na None, aby sygnalizować problem
 
 
 @app.get("/")
@@ -54,12 +58,12 @@ async def root():
     """
     return {"message": "Bielik Text Generation API is running!"}
 
+
 @app.post("/generate")
 async def generate_text(request: GenerationRequest):
     """
     Endpoint do generowania tekstu na podstawie promptu.
     """
-    print(request)
     if generator is None:
         raise HTTPException(status_code=503, detail="Model nie został załadowany lub wystąpił błąd.")
 
@@ -69,13 +73,16 @@ async def generate_text(request: GenerationRequest):
             max_new_tokens=request.max_new_tokens,
             temperature=request.temperature,
             top_p=request.top_p,
-            do_sample=True, # Ważne dla generowania z temperaturą
+            do_sample=True,  # Ważne dla generowania z temperaturą
         )
         # Pipeline zwraca listę słowników, bierzemy pierwszy wynik
-        return {"generated_text": generated_text[0]["generated_text"]}
+        #gen_text = {"generated_text": generated_text[0]["generated_text"]}
+        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)))
+    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)))