GrootW/newspace

1from fastapi import FastAPI, HTTPException2from pydantic import BaseModel3from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline4import torch5import numpy as np # Проверка версии NumPy6 7app = FastAPI()8 9 10# Проверка версии NumPy11assert np.__version__.startswith('1.'), f"Несовместимая версия NumPy: {np.__version__}"12 13 14 15class RequestData(BaseModel):16    prompt: str17    max_tokens: int = 5018 19MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"20 21try:22    # Загрузка модели с явным указанием device_map23    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)24    model = AutoModelForCausalLM.from_pretrained(25        MODEL_NAME,26        torch_dtype=torch.float32,27        device_map="auto",28        low_cpu_mem_usage=True29    )30    31    # Создаем pipeline без указания device32    generator = pipeline(33        "text-generation",34        model=model,35        tokenizer=tokenizer36    )37except Exception as e:38    print(f"Ошибка загрузки модели: {str(e)}")39    generator = None40 41@app.post("/generate")42async def generate_text(request: RequestData):43    if not generator:44        raise HTTPException(status_code=503, detail="Модель не загружена")45    46    try:47        output = generator(48            request.prompt,49            max_new_tokens=request.max_tokens,50            do_sample=False,51            num_beams=1,52            temperature=0.7,53        )54        return {"response": output[0]["generated_text"]}55    except Exception as e:56        raise HTTPException(status_code=500, detail=str(e))57 58@app.get("/health")59async def health_check():60    return {"status": "ok" if generator else "unavailable"}