Spaces:
Runtime error
Runtime error
File size: 916 Bytes
5d96d6a b4e99db c3325d4 414f5f4 b4e99db 5d96d6a 2ba48b4 5fb8929 5d96d6a b4e99db 5fb8929 414f5f4 b4e99db b96ee8b b4e99db b96ee8b b4e99db 5d96d6a c1e79d1 |
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 |
from fastapi import FastAPI, Request
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
app = FastAPI()
model_id = "google/flan-t5-small" # Replace with your model here
#"unsloth/mistral-7b-v0.2-bnb-4bit"
#deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.float16,
device_map="auto",
)
@app.post("/generate")
async def generate(request: Request):
data = await request.json()
prompt = data.get("prompt", "").strip()
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(
inputs.input_ids,
max_new_tokens=100,
use_cache=True,
temperature=0.7,
)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
return {"output": generated_text}
|