Update app.py

app.py

@@ -1,39 +1,61 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import time
+from transformers import LlamaForCausalLM, PreTrainedTokenizerFast, pipeline
 
-model_name = "Azurro/APT-1B-Base"
+model_name = "Azurro/APT3-1B-Instruct-v1"
 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+model = LlamaForCausalLM.from_pretrained(model_name, torch.float16)
 
-generator = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-)
-
-def generate_text(prompt, max_length, temperature, top_k, top_p, beams):
-    output = generator(prompt, 
-                       max_length=max_length, 
-                       temperature=temperature, 
-                       top_k=top_k,
-                       do_sample=True,
-                       top_p=top_p,
-                       num_beams=beams)
-    return output[0]['generated_text']
-
-input_text = gr.inputs.Textbox(label="Input Text")
-max_length = gr.inputs.Slider(1, 100, step=1, default=30, label="Max Length")
-temperature = gr.inputs.Slider(0.1, 1.0, step=0.1, default=0.8, label="Temperature")
-top_k = gr.inputs.Slider(1, 200, step=1, default=10, label="Top K")
-top_p = gr.inputs.Slider(0.1, 2.0, step=0.1, default=0.95, label="Top P")
-beams = gr.inputs.Slider(1, 20, step=1, default=1, label="Beams")
+def generate_text(prompt, max_length, temperature, top_k, top_p):
+    prompt = f'<s>[INST] {prompt.strip()} [/INST]'
+    input_ids = tokenizer(prompt, return_tensors='pt', add_special_tokens=False).input_ids.to(model.device)
+    start_time = time.time()
+    output = model.generate(
+        inputs=input_ids, 
+        max_new_tokens=max_length,
+        temperature=temperature, 
+        top_k=top_k,
+        do_sample=(temperature > 0),
+        top_p=top_p,
+        num_beams=1,
+        bos_token_id=1,
+        eos_token_id=2,
+        pad_token_id=3,
+        repetition_penalty=1.1
+    )
+    elapsed_time = time.time() - start_time
+    decoded_output = tokenizer.decode(output[0])
+    input_tokens_count = len(input_ids[0])
+    input_chars_count = len(prompt)
+    output_tokens_count = len(output[0])
+    output_chars_count = len(decoded_output)
+    gen_speed = output_tokens_count / elapsed_time
+    decoded_output = decoded_output[len(prompt):].replace('</s>','').strip()
+    print(f"Input tokens: {input_tokens_count} (chars: {input_chars_count}), Output tokens: {output_tokens_count} (chars: {output_chars_count}), Gen Time: {elapsed_time:.2f} secs ({gen_speed} toks/sec)")
+    print(f"{'*'*10} Input {'*'*10}\n{prompt}")
+    print(f"{'*'*10} Output {'*'*10}\n{prompt}")
+    print(f"{'*'*30}")
+    return decoded_output, input_tokens_count, input_chars_count, output_tokens_count, output_chars_count, gen_speed
 
-outputs = gr.outputs.Textbox(label="Generated Text")
-
-iface = gr.Interface(generate_text, inputs=[input_text, max_length, temperature, top_k, top_p, beams], outputs=outputs)
-iface.queue(concurrency_count=1)
-iface.launch(max_threads=100)
+demo = gr.Interface(
+    fn=generate_text,
+    inputs=[
+        gr.inputs.Textbox(label="Input Text"), 
+        gr.inputs.Slider(1, 1000, step=1, default=100, label="Max Length"), 
+        gr.inputs.Slider(0.0, 1.5, step=0.1, default=0.6, label="Temperature"), 
+        gr.inputs.Slider(1, 400, step=1, default=200, label="Top K"), 
+        gr.inputs.Slider(0.0, 1.0, step=0.05, default=0.95, label="Top P")
+    ],
+    outputs=[
+        gr.outputs.Textbox(label="Generated Text"),
+        gr.outputs.Textbox(label="Input Tokens Count"),
+        gr.outputs.Textbox(label="Input Characters Count"),
+        gr.outputs.Textbox(label="Output Tokens Count"),
+        gr.outputs.Textbox(label="Output Characters Count"),
+        gr.outputs.Textbox(label="Generation speed in tokens per second"),
+    ]
+)
+demo.queue(concurrency_count=1)
+demo.launch(max_threads=20)