rishiraj commited on
Commit
9157017
·
verified ·
1 Parent(s): 946ab7a

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +131 -0
app.py ADDED
@@ -0,0 +1,131 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from threading import Thread
3
+ from typing import Iterator
4
+
5
+ import gradio as gr
6
+ import spaces
7
+ import torch
8
+ from transformers import AutoModelForCausalLM, GemmaTokenizerFast, TextIteratorStreamer
9
+
10
+ DESCRIPTION = """\
11
+ # Prompt Generator with Gemma 2 9B IT
12
+ Sometimes, the hardest part of using an AI model is figuring out how to prompt it effectively. To help with this, we’ve created a prompt generation tool that guides Gemma to generate high-quality prompt templates tailored to your specific tasks. These templates follow many of our prompt engineering best practices.
13
+
14
+ The prompt generator is particularly useful as a tool for solving the “blank page problem” to give you a jumping-off point for further testing and iteration.
15
+ """
16
+
17
+ MAX_MAX_NEW_TOKENS = 2048
18
+ DEFAULT_MAX_NEW_TOKENS = 1024
19
+ MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
20
+
21
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
22
+
23
+ model_id = "unsloth/gemma-2-9b-it"
24
+ tokenizer = GemmaTokenizerFast.from_pretrained(model_id)
25
+ model = AutoModelForCausalLM.from_pretrained(
26
+ model_id,
27
+ device_map="auto",
28
+ torch_dtype=torch.bfloat16,
29
+ )
30
+ model.config.sliding_window = 4096
31
+ model.eval()
32
+
33
+
34
+ @spaces.GPU(duration=90)
35
+ def generate(
36
+ message: str,
37
+ chat_history: list[dict],
38
+ max_new_tokens: int = 1024,
39
+ temperature: float = 0.6,
40
+ top_p: float = 0.9,
41
+ top_k: int = 50,
42
+ repetition_penalty: float = 1.2,
43
+ ) -> Iterator[str]:
44
+ conversation = chat_history.copy()
45
+ conversation.append({"role": "user", "content": message})
46
+
47
+ input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
48
+ if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
49
+ input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
50
+ gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
51
+ input_ids = input_ids.to(model.device)
52
+
53
+ streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
54
+ generate_kwargs = dict(
55
+ {"input_ids": input_ids},
56
+ streamer=streamer,
57
+ max_new_tokens=max_new_tokens,
58
+ do_sample=True,
59
+ top_p=top_p,
60
+ top_k=top_k,
61
+ temperature=temperature,
62
+ num_beams=1,
63
+ repetition_penalty=repetition_penalty,
64
+ )
65
+ t = Thread(target=model.generate, kwargs=generate_kwargs)
66
+ t.start()
67
+
68
+ outputs = []
69
+ for text in streamer:
70
+ outputs.append(text)
71
+ yield "".join(outputs)
72
+
73
+
74
+ chat_interface = gr.ChatInterface(
75
+ fn=generate,
76
+ additional_inputs=[
77
+ gr.Slider(
78
+ label="Max new tokens",
79
+ minimum=1,
80
+ maximum=MAX_MAX_NEW_TOKENS,
81
+ step=1,
82
+ value=DEFAULT_MAX_NEW_TOKENS,
83
+ ),
84
+ gr.Slider(
85
+ label="Temperature",
86
+ minimum=0.1,
87
+ maximum=4.0,
88
+ step=0.1,
89
+ value=0.6,
90
+ ),
91
+ gr.Slider(
92
+ label="Top-p (nucleus sampling)",
93
+ minimum=0.05,
94
+ maximum=1.0,
95
+ step=0.05,
96
+ value=0.9,
97
+ ),
98
+ gr.Slider(
99
+ label="Top-k",
100
+ minimum=1,
101
+ maximum=1000,
102
+ step=1,
103
+ value=50,
104
+ ),
105
+ gr.Slider(
106
+ label="Repetition penalty",
107
+ minimum=1.0,
108
+ maximum=2.0,
109
+ step=0.05,
110
+ value=1.2,
111
+ ),
112
+ ],
113
+ stop_btn=None,
114
+ examples=[
115
+ ["Hello there! How are you doing?"],
116
+ ["Can you explain briefly to me what is the Python programming language?"],
117
+ ["Explain the plot of Cinderella in a sentence."],
118
+ ["How many hours does it take a man to eat a Helicopter?"],
119
+ ["Write a 100-word article on 'Benefits of Open-Source in AI research'"],
120
+ ],
121
+ cache_examples=False,
122
+ type="messages",
123
+ )
124
+
125
+ with gr.Blocks(css="style.css", fill_height=True) as demo:
126
+ gr.Markdown(DESCRIPTION)
127
+ gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
128
+ chat_interface.render()
129
+
130
+ if __name__ == "__main__":
131
+ demo.queue(max_size=20).launch()