Spaces:
Paused
Paused
File size: 1,340 Bytes
4468cfe 45123df 4468cfe e05b36f 4468cfe e05b36f 4468cfe 45123df 9d98b84 45123df |
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 |
from fastapi import FastAPI, Request
from transformers import AutoModelForCausalLM, AutoTokenizer
from fastapi.responses import StreamingResponse
import torch
app = FastAPI()
# Load the model and tokenizer
model_name = "EleutherAI/gpt-neo-1.3B" # Replace with your desired model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
@app.get("/")
def read_root():
return {"Hello": "World"}
@app.post("/predict")
async def predict(request: Request):
data = await request.json()
prompt = data.get("prompt", "")
if not prompt:
return {"error": "Prompt is required"}
# Tokenize the input
inputs = tokenizer(prompt, return_tensors="pt").to("cpu") # Use "cuda" if GPU is enabled
# Generator function to stream tokens
def token_generator():
outputs = model.generate(
inputs.input_ids,
max_length=40,
do_sample=True,
num_return_sequences=1,
temperature=0.7,
top_p=0.9,
repetition_penalty=1.2,
)
for token_id in outputs[0]:
token = tokenizer.decode(token_id, skip_special_tokens=True)
yield f"{token} "
# Return StreamingResponse
return StreamingResponse(token_generator(), media_type="text/plain")
|