Update app.py

app.py

@@ -1,72 +1,49 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import time
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
+from threading import Thread
 
 # Load model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T")
 model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T")
 
-def generate_text(prompt, max_length=100, min_length=20, temperature=1.0):
-    # Tokenize the prompt
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
-
-    # Generate text
-    output = model.generate(
-        input_ids,
-        max_length=max_length,
-        min_length=min_length,
-        num_return_sequences=1,
-        temperature=temperature
-    )
-
-    # Decode the generated output
-    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
-
-    return generated_text
-
-def print_like_dislike(x: gr.LikeData):
-    print(x.index, x.value, x.liked)
-
-def add_text(history, text):
-    history = history + [(text, None)]
-    return history, gr.Textbox(value="", interactive=False)
-
-def bot(history, max_length=100, min_length=20, temperature=1.0):
-    prompt = history[-1][0]
-    response = generate_text(prompt, max_length=max_length, min_length=min_length, temperature=temperature)
-    history[-1][1] = ""
-    for character in response:
-        history[-1][1] += character
-        time.sleep(0.05)
-        yield history
-
-with gr.Blocks() as demo:
-    chatbot = gr.Chatbot(
-        [],
-        elem_id="chatbot",
-        bubble_full_width=False,
-        avatar_images=(None, None),  # Set avatar image path or URL
-    )
-
-    with gr.Row():
-        txt = gr.Textbox(
-            scale=4,
-            show_label=False,
-            placeholder="Enter text and press enter, or upload an image",
-            container=False,
+class StopOnTokens(StoppingCriteria):
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        stop_ids = [29, 0]
+        for stop_id in stop_ids:
+            if input_ids[0][-1] == stop_id:
+                return True
+        return False
+
+def predict(message, history):
+
+    history_transformer_format = history + [[message, ""]]
+    stop = StopOnTokens()
+
+    messages = "".join(["".join(["\n<human>:"+item[0], "\n<bot>:"+item[1]])  #curr_system_message +
+                for item in history_transformer_format])
+
+    model_inputs = tokenizer([messages], return_tensors="pt").to("cuda")
+    streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
+    generate_kwargs = dict(
+        model_inputs,
+        streamer=streamer,
+        max_new_tokens=1024,
+        do_sample=True,
+        top_p=1,
+        top_k=50,
+        temperature=1.0,
+        num_beams=1,
+        stopping_criteria=StoppingCriteriaList([stop])
         )
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
 
-        max_len_slider = gr.Slider(0, 2048, 100, label="Max Length")
-        min_len_slider = gr.Slider(0, 2048, 20, label="Min Length")
-        temp_slider = gr.Slider(0.1, 2.0, 1.0, label="Temperature")
-
-    txt_msg = txt.submit(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
-        bot, chatbot, chatbot, max_len_slider, min_len_slider, temp_slider
-    )
-    txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)
+    partial_message  = ""
+    for new_token in streamer:
+        if new_token != '<':
+            partial_message += new_token
+            yield partial_message
 
-    chatbot.like(print_like_dislike, None, None)
 
-demo.queue()
-if __name__ == "__main__":
-    demo.launch()
+gr.ChatInterface(predict).queue().launch()