Tonic commited on
Commit
709d394
1 Parent(s): 62a2be1

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +43 -0
app.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import spaces
2
+ import gradio as gr
3
+ from transformers import AutoTokenizer, AutoModelForCausalLM
4
+ import torch
5
+
6
+ # Load the tokenizer and model
7
+ model_id = "CohereForAI/c4ai-command-r-v01"
8
+ tokenizer = AutoTokenizer.from_pretrained(model_id)
9
+ model = AutoModelForCausalLM.from_pretrained(model_id)
10
+
11
+ @spaces.GPU
12
+ def generate_response(user_input, max_new_tokens, temperature):
13
+ # Format message with the command-r chat template
14
+ messages = [{"role": "user", "content": user_input}]
15
+ input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
16
+
17
+ # Generate tokens
18
+ gen_tokens = model.generate(
19
+ input_ids['input_ids'],
20
+ max_length=max_new_tokens + input_ids['input_ids'].shape[1], # Adjusting max_length to account for input length
21
+ do_sample=True,
22
+ temperature=temperature,
23
+ )
24
+
25
+ # Decode tokens to string
26
+ gen_text = tokenizer.decode(gen_tokens[0])
27
+ return gen_text
28
+
29
+ # Define the Gradio interface
30
+ iface = gr.Interface(
31
+ fn=generate_response,
32
+ inputs=[
33
+ gr.inputs.Textbox(lines=2, label="Your Message"),
34
+ gr.inputs.Slider(minimum=10, maximum=100, default=50, label="Max New Tokens"),
35
+ gr.inputs.Slider(minimum=0.1, maximum=1.0, step=0.1, default=0.3, label="Temperature")
36
+ ],
37
+ outputs=gr.outputs.Textbox(label="Model Response"),
38
+ title="Text Generation Model Interface",
39
+ description="This is a Gradio interface for a text generation model. Enter your message and adjust the parameters to generate a response."
40
+ )
41
+
42
+ # Launch the application
43
+ iface.launch()