|
|
from fastapi import FastAPI, Request
|
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
|
|
|
import torch
|
|
|
|
|
|
app = FastAPI()
|
|
|
|
|
|
model_name = "howtomakepplragequit/phi2-lora-instruct"
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
|
model_name,
|
|
|
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
|
|
|
device_map="auto"
|
|
|
)
|
|
|
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
|
|
|
|
|
|
@app.post("/generate")
|
|
|
async def generate(request: Request):
|
|
|
data = await request.json()
|
|
|
prompt = data.get("prompt", "")
|
|
|
formatted = f"### Instruction:\n{prompt}\n\n### Response:\n"
|
|
|
result = pipe(formatted, max_new_tokens=200)[0]["generated_text"]
|
|
|
return {"response": result.split("### Response:")[-1].strip()}
|
|
|
|
