Spaces:
Runtime error
Runtime error
File size: 892 Bytes
9df6db6 1c3cbe0 de271c2 9df6db6 de271c2 9df6db6 1c3cbe0 de271c2 58c3fbf de271c2 58c3fbf 9df6db6 4303ef8 e10efef 1c3cbe0 de271c2 |
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 |
from fastapi import FastAPI
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
model_id = "meta-llama/Meta-Llama-Guard-2-8B"
device = "cuda"
dtype = torch.bfloat16
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype, device_map=device)
app = FastAPI()
def moderate(chat):
input_ids = tokenizer.apply_chat_template(chat, return_tensors="pt").to(device)
output = model.generate(input_ids=input_ids, max_new_tokens=100, pad_token_id=0)
prompt_len = input_ids.shape[-1]
return tokenizer.decode(output[0][prompt_len:], skip_special_tokens=True)
@app.get("/")
def greet_json():
return {"Hello": "World!"}
@app.post("/name")
def helloName(input):
result = moderate([
{"role": "user", "content": "()".format(input)}
])
return {"hello": "()".format(result)}
|