|
|
from fastapi import FastAPI, Request |
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM |
|
|
|
|
|
app = FastAPI() |
|
|
|
|
|
model_path = "tinyllama-finetuned-skin" |
|
|
tokenizer = AutoTokenizer.from_pretrained(model_path) |
|
|
model = AutoModelForCausalLM.from_pretrained(model_path) |
|
|
|
|
|
@app.post("/generate") |
|
|
async def generate(request: Request): |
|
|
data = await request.json() |
|
|
prompt = data["prompt"] |
|
|
inputs = tokenizer(prompt, return_tensors="pt") |
|
|
outputs = model.generate(**inputs, max_new_tokens=100) |
|
|
result = tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
|
return {"response": result} |
