Spaces:
Running
Running
File size: 1,479 Bytes
dea3a07 2fc7e1b 115a37b d1fd8de 115a37b 2fc7e1b dea3a07 2fc7e1b e604a26 861971b af0df21 d1fd8de 2fc7e1b af0df21 2fc7e1b d1fd8de 2fc7e1b d1fd8de 2fc7e1b e604a26 115a37b 2fc7e1b 115a37b af0df21 2fc7e1b 115a37b 2fc7e1b 115a37b |
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 37 38 39 40 41 42 43 44 45 46 47 48 |
import os
import torch
from fastapi import FastAPI
from transformers import AutoTokenizer, AutoModelForCausalLM
from pydantic import BaseModel
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
app = FastAPI()
model_name = "google/gemma-2-2b-it"
tokenizer = None
model = None
try:
logger.info(f"Loading model: {model_name}")
tokenizer = AutoTokenizer.from_pretrained(model_name, token=os.getenv("HF_TOKEN"))
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16, # メモリ削減
device_map="cpu", # GPU利用不可
token=os.getenv("HF_TOKEN"),
low_cpu_mem_usage=True
)
logger.info("Model loaded successfully")
except Exception as e:
logger.error(f"Model load error: {e}")
raise
class TextInput(BaseModel):
text: str
max_length: int = 50
@app.post("/generate")
async def generate_text(input: TextInput):
try:
logger.info(f"Generating text for input: {input.text}")
inputs = tokenizer(input.text, return_tensors="pt", max_length=512, truncation=True).to("cpu")
outputs = model.generate(**inputs, max_length=input.max_length)
result = tokenizer.decode(outputs[0], skip_special_tokens=True)
logger.info(f"Generated text: {result}")
return {"generated_text": result}
except Exception as e:
logger.error(f"Generation error: {e}")
return {"error": str(e)}"
|