forked from ShishirPatil/gorilla
-
Notifications
You must be signed in to change notification settings - Fork 0
/
predict.py
64 lines (54 loc) · 2.02 KB
/
predict.py
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
import os
from cog import BasePredictor, Input
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
def get_prompt(user_query: str) -> str:
"""
Generates a conversation prompt based on the user's query.
Parameters:
- user_query (str): The user's query.
Returns:
- str: The formatted conversation prompt.
"""
return f"USER: <<question>> {user_query}\nASSISTANT: "
class Predictor(BasePredictor):
def setup(self):
"""
Load the model into memory to make running multiple predictions efficient.
Sets up the device, model, tokenizer, and pipeline for text generation.
"""
# Device setup
self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
self.torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
# Model and tokenizer setup
model_id = os.environ.get('MODEL', 'gorilla-openfunctions-v2')
self.tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
self.model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=self.torch_dtype,
low_cpu_mem_usage=True,
trust_remote_code=True
)
# Move model to device
self.model.to(self.device)
# Pipeline setup
self.pipe = pipeline(
"text-generation",
model=self.model,
tokenizer=self.tokenizer,
max_new_tokens=256,
batch_size=16,
torch_dtype=self.torch_dtype,
device=self.device,
)
def predict(self, user_query: str = Input(description="User's query")) -> str:
"""
Run a single prediction on the model using the provided user query.
Parameters:
- user_query (str): The user's query for the model.
Returns:
- str: The model's generated text based on the query.
"""
prompt = get_prompt(user_query)
output = self.pipe(prompt)
return output