Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -4,12 +4,14 @@ import spaces
|
|
4 |
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
|
5 |
|
6 |
|
7 |
-
model_name = "
|
8 |
token = os.environ['hf_token']
|
9 |
|
10 |
pipe = pipeline("text-generation", model=model_name, device="cuda")
|
11 |
|
12 |
|
|
|
|
|
13 |
system_prompt = '''You are given an input text for a chat interface. Propose auto-completion to the text. You have several roles:
|
14 |
- Fight under-specification: if the user does not provide sufficient context, propose them a set of relevant suggestions.
|
15 |
- Complete text: The text provided to you is in the making. If you have a good idea for how to complete - make suggestions.
|
@@ -24,7 +26,7 @@ def generate(text):
|
|
24 |
{'role': 'system', 'content': system_prompt},
|
25 |
{'role': 'user', 'content': text}
|
26 |
]
|
27 |
-
return pipe(messages)
|
28 |
|
29 |
|
30 |
if __name__ == "__main__":
|
|
|
4 |
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
|
5 |
|
6 |
|
7 |
+
model_name = "TheBloke/OpenHermes-2.5-Mistral-7B-GGUF"
|
8 |
token = os.environ['hf_token']
|
9 |
|
10 |
pipe = pipeline("text-generation", model=model_name, device="cuda")
|
11 |
|
12 |
|
13 |
+
generate_kwargs = {'max_new_tokens': 20}
|
14 |
+
|
15 |
system_prompt = '''You are given an input text for a chat interface. Propose auto-completion to the text. You have several roles:
|
16 |
- Fight under-specification: if the user does not provide sufficient context, propose them a set of relevant suggestions.
|
17 |
- Complete text: The text provided to you is in the making. If you have a good idea for how to complete - make suggestions.
|
|
|
26 |
{'role': 'system', 'content': system_prompt},
|
27 |
{'role': 'user', 'content': text}
|
28 |
]
|
29 |
+
return pipe(messages, generate_kwargs=generate_kwargs)
|
30 |
|
31 |
|
32 |
if __name__ == "__main__":
|