EndorAI commited on
Commit
a124f60
1 Parent(s): 2bd0164

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +33 -0
app.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList
3
+
4
+ tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-tuned-alpha-7b")
5
+ model = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-tuned-alpha-7b")
6
+ model.half().cuda()
7
+
8
+ class StopOnTokens(StoppingCriteria):
9
+ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
10
+ stop_ids = [50278, 50279, 50277, 1, 0]
11
+ for stop_id in stop_ids:
12
+ if input_ids[0][-1] == stop_id:
13
+ return True
14
+ return False
15
+
16
+ system_prompt = """<|SYSTEM|># StableLM Tuned (Alpha version)
17
+ - StableLM is a helpful and harmless open-source AI language model developed by StabilityAI.
18
+ - StableLM is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.
19
+ - StableLM is more than just an information source, StableLM is also able to write poetry, short stories, and make jokes.
20
+ - StableLM will refuse to participate in anything that could harm a human.
21
+ """
22
+
23
+ prompt = f"{system_prompt}<|USER|>What's your mood today?<|ASSISTANT|>"
24
+
25
+ inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
26
+ tokens = model.generate(
27
+ **inputs,
28
+ max_new_tokens=64,
29
+ temperature=0.7,
30
+ do_sample=True,
31
+ stopping_criteria=StoppingCriteriaList([StopOnTokens()])
32
+ )
33
+ print(tokenizer.decode(tokens[0], skip_special_tokens=True))