Tonic commited on
Commit
651f3a2
β€’
1 Parent(s): 6d32964

add minitron 8B base

Browse files
Files changed (2) hide show
  1. app.py +41 -5
  2. transformers +1 -0
app.py CHANGED
@@ -1,7 +1,43 @@
1
- import gradio as gr
 
 
2
 
3
- def greet(name):
4
- return "Hello " + name + "!!"
5
 
6
- demo = gr.Interface(fn=greet, inputs="text", outputs="text")
7
- demo.launch()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from transformers import AutoTokenizer, AutoModelForCausalLM
4
 
5
+ title = """# πŸ™‹πŸ»β€β™‚οΈ Welcome to Tonic's Minitron-8B-Base"""
 
6
 
7
+ # Load the tokenizer and model
8
+ model_path = "nvidia/Minitron-8B-Base"
9
+ tokenizer = AutoTokenizer.from_pretrained(model_path)
10
+
11
+ device='cuda'
12
+ dtype=torch.bfloat16
13
+ model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=dtype, device_map=device)
14
+
15
+ # Define the prompt format
16
+ def create_prompt(instruction):
17
+ PROMPT = '''You are TronTonic an AI created by Tonic-AI. Below is an instruction that describes a task.\n\nWrite a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Response:'''
18
+ return PROMPT.format(instruction=instruction)
19
+
20
+ def respond(message, history, system_message, max_tokens, temperature, top_p):
21
+ prompt = create_prompt(message)
22
+
23
+ input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
24
+
25
+ output_ids = model.generate(input_ids, max_length=50, num_return_sequences=1)
26
+
27
+ output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
28
+
29
+ return output_text
30
+
31
+ demo = gr.ChatInterface(
32
+ gr.markdown(title),
33
+ # gr.markdown(description),
34
+ respond,
35
+ additional_inputs=[
36
+ gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
37
+ gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
38
+ gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")
39
+ ],
40
+ )
41
+
42
+ if __name__ == "__main__":
43
+ demo.launch()
transformers ADDED
@@ -0,0 +1 @@
 
 
1
+ Subproject commit 63d9cb0afd2bf5d4cb5431ba1b2c4e353752a937