Update main.py
Browse files
main.py
CHANGED
@@ -6,12 +6,13 @@ import prompt_style
|
|
6 |
|
7 |
|
8 |
model_id = "failspy/Meta-Llama-3-8B-Instruct-abliterated-v3"
|
9 |
-
client =
|
10 |
|
11 |
class Item(BaseModel):
|
12 |
prompt: str
|
13 |
history: list
|
14 |
system_prompt: str
|
|
|
15 |
temperature: float = 0.6
|
16 |
max_new_tokens: int = 1024
|
17 |
top_p: float = 0.95
|
@@ -44,6 +45,8 @@ def generate(item: Item):
|
|
44 |
seed=item.seed,
|
45 |
)
|
46 |
|
|
|
|
|
47 |
formatted_prompt = format_prompt(item)
|
48 |
stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
|
49 |
output = ""
|
|
|
6 |
|
7 |
|
8 |
model_id = "failspy/Meta-Llama-3-8B-Instruct-abliterated-v3"
|
9 |
+
client = None
|
10 |
|
11 |
class Item(BaseModel):
|
12 |
prompt: str
|
13 |
history: list
|
14 |
system_prompt: str
|
15 |
+
token:str
|
16 |
temperature: float = 0.6
|
17 |
max_new_tokens: int = 1024
|
18 |
top_p: float = 0.95
|
|
|
45 |
seed=item.seed,
|
46 |
)
|
47 |
|
48 |
+
if not client:
|
49 |
+
client = InferenceClient(token=item.token, model=model_id)
|
50 |
formatted_prompt = format_prompt(item)
|
51 |
stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
|
52 |
output = ""
|