File size: 1,159 Bytes
de9ab11
92b4185
ed0a7a7
de9ab11
cee5832
92b4185
 
de9ab11
 
923d7e0
de9ab11
923d7e0
4e5b21c
 
 
 
de9ab11
4e5b21c
 
 
 
 
de9ab11
 
92b4185
 
 
 
 
4e5b21c
 
 
 
 
de9ab11
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
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))