Spaces:
Running
Running
File size: 5,250 Bytes
38826eb |
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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 |
import argparse
import os
import time
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Dict, List, Optional
import torch
import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.testclient import TestClient
from transformers import AutoModelForCausalLM, AutoTokenizer
from custom_llm_inference import get_highlights_inner, get_next_token_predictions_inner
ml_models = {}
parser = argparse.ArgumentParser()
parser.add_argument("--gpu", action="store_true", help="Enable GPU usage")
args = parser.parse_args()
USE_GPU = args.gpu
if not USE_GPU:
print("Running without GPU. To enable GPU, run with the --gpu flag.")
@asynccontextmanager
async def models_lifespan(app: FastAPI):
#model_name = 'google/gemma-1.1-7b-it'
#model_name = 'google/gemma-1.1-2b-it'
model_name = 'google/gemma-2-9b-it'
dtype = torch.bfloat16 if USE_GPU else torch.float16
ml_models["llm"] = llm = {
'tokenizer': AutoTokenizer.from_pretrained(model_name),
'model': AutoModelForCausalLM.from_pretrained(model_name, device_map="auto" if USE_GPU else "cpu", torch_dtype=dtype)
}
print("Loaded llm with device map:")
print(llm['model'].hf_device_map)
# Print timing info for each endpoint
print("\nRunning endpoint tests...")
test_doc = "This is a test document that needs to be revised for clarity and conciseness."
test_prompt = "Make this more clear and concise."
client = TestClient(app)
start = time.time()
response = client.get("/api/highlights",
params={"doc": test_doc, "prompt": test_prompt})
print(f"Highlights endpoint: {time.time() - start:.2f}s")
start = time.time()
response = client.get("/api/next_token",
params={"original_doc": test_doc, "prompt": test_prompt, "doc_in_progress": "This is"})
print(f"Next token endpoint: {time.time() - start:.2f}s")
start = time.time()
response = client.get("/api/gen_revisions",
params={"doc": test_doc, "prompt": test_prompt, "n": 1})
print(f"Gen revisions endpoint: {time.time() - start:.2f}s")
yield
# Release resources on exit
ml_models.clear()
DEBUG = os.getenv("DEBUG") or False
PORT = int(os.getenv("PORT") or "19570")
app = FastAPI(lifespan=models_lifespan)
origins = [
"*",
]
app.add_middleware(
CORSMiddleware,
allow_origins=origins,
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
@app.get("/api/highlights")
def get_highlights(doc: str, prompt: Optional[str] = None, updated_doc: Optional[str] = '', k: Optional[int] = 5):
''' Example of using this in JavaScript:
let url = new URL('http://localhost:8000/api/highlights')
url.searchParams.append('doc', 'This is a test document. It is a test document because it is a test document.')
url.searchParams.append('prompt', 'Rewrite this document to be more concise.')
url.searchParams.append('updated_doc', 'This is a test document.')
let response = await fetch(url)
'''
llm = ml_models['llm']
model = llm['model']
tokenizer = llm['tokenizer']
if prompt is None:
prompt = "Rewrite this document to be more concise."
highlights = get_highlights_inner(model, tokenizer, doc, prompt, updated_doc, k)
return {'highlights': highlights}
@app.get('/api/next_token')
def get_next_token_predictions(original_doc: str,
prompt: str,
doc_in_progress: str,
k: Optional[int] = 5):
model = ml_models['llm']['model']
tokenizer = ml_models['llm']['tokenizer']
decoded_next_tokens, next_token_logits = get_next_token_predictions_inner(
model, tokenizer, original_doc, prompt, doc_in_progress, k)
return {
'next_tokens': decoded_next_tokens
}
@app.get('/api/gen_revisions')
def gen_revisions(
prompt: str,
doc: str,
n: Optional[int] = 5):
model = ml_models['llm']['model']
tokenizer = ml_models['llm']['tokenizer']
messages = [
{
"role": "user",
"content": f"{prompt}\n\n{doc}",
},
]
tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)
generations = model.generate(
tokenized_chat, num_return_sequences=n,
max_length=1024, do_sample=True, top_k=50, top_p=0.95, temperature=0.5,
return_dict_in_generate=True, output_scores=True)
generated_docs = tokenizer.batch_decode(generations.sequences, skip_special_tokens=True)
#print(generations.scores)
# Remove prompt text. see https://github.com/huggingface/transformers/blob/v4.46.2/src/transformers/pipelines/text_generation.py#L37
prompt_length = len(
tokenizer.decode(
tokenized_chat[0],
skip_special_tokens=True,
clean_up_tokenization_spaces=True,
))
return {
'revised_docs': [dict(doc_text=doc[prompt_length:]) for doc in generated_docs]
}
if __name__ == "__main__":
uvicorn.run(app, host="localhost", port=PORT)
|