Spaces:
Runtime error
Runtime error
File size: 3,651 Bytes
9131d62 dff38c9 9131d62 7df77f3 dff38c9 9131d62 dff38c9 9131d62 afc30f3 9131d62 dff38c9 9131d62 afc30f3 9131d62 afc30f3 dff38c9 afc30f3 9131d62 afc30f3 | 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 | import os
import requests
import time
# from google.cloud import storage
from sentence_transformers import SentenceTransformer
from .config import max_new_tokens, streaming_url, job_url, default_payload, headers, embedding_path
from .db.db_utilities import query_db
class Model:
'''Client class for holding Llama2 model and tokenizer. Models are loaded according to
ENVIRONMENT environment variable
'''
def __init__(self,
max_new_tokens:int=max_new_tokens):
self.max_new_tokens = max_new_tokens
# self.embedding_model = SentenceTransformer('multi-qa-mpnet-base-dot-v1')
self.embedding_model = SentenceTransformer(embedding_path)
def inference(self, query:str, table:str):
'''Inference function for gradio text streaming'''
# set in the case that None
if table == 'Court' or table is None:
table= 'court_opinion'
output= [[None, '']]
for i in self.query_model(query, table):
output[0][1] += i
yield output
def get_context(self, query:str, table:str='court_opinion')-> str:
"""Query vectordb for additional context and compiles a new query string with added context"""
matches = query_db(query, self.embedding_model, table=table)
if len(matches) > 0:
match = '"""' + matches[0][0] + '"""'
context = "You are the United States Supreme Court. Use the following historical opinion to give your ruling on a court case description. Historical opinion: " + match
else:
context = 'You are the United States Supreme Court. Give your ruling on a court case description.'
return context + " Answer in less than 400 words in the format Opinion: <opinion> "
def query_model(self, query:str, table:str, default_payload:dict=default_payload, timeout:int=60, **kwargs) -> str:
"""Query the model api on runpod. Runs for 60s by default. Generator response until job is complete"""
context = self.get_context(query=query, table=table)
for k,v in kwargs:
default_payload['input']['sampling_params'][k] = v
augmented_prompt_template = [
{
"role": "system",
"content": context,
},
{
"role": "user",
"content": query,
}
]
default_payload["input"]["prompt"] = augmented_prompt_template
job_id = requests.post(job_url, json=default_payload, headers=headers).json()['id']
for i in range(timeout):
time.sleep(1)
stream_response = requests.get(streaming_url+ job_id, headers=headers).json()
if stream_response['status'] == 'COMPLETED':
break
for i in stream_response['stream']:
for j in i['output']['text']:
yield j
# def download_checkpoints(self, bucket_name: str = bucket_name):
# """Downloads model files from gcp storage if running in gcp."""
# if not(os.path.exists('model/')):
# os.mkdir('model')
# storage_client = storage.Client()
# bucket = storage_client.bucket(bucket_name)
# # get tokenizer
# blob = bucket.blob(self.tokenizer_path)
# blob.download_to_filename(self.tokenizer_path)
# # get model files to models/
# model_file_paths = [self.model_path + i for i in model_files]
# for object_name in model_file_paths:
# blob = bucket.blob(object_name)
# blob.download_to_filename(object_name)
|