update

app.py

@@ -16,6 +16,7 @@ repo_id = "j2moreno/test-model"
 
 model = AutoModelForCausalLM.from_pretrained(repo_id).to(device)
 tokenizer = AutoTokenizer.from_pretrained(repo_id)
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 SEED = 42
 
@@ -43,24 +44,47 @@ examples=[
 #         print("You downvoted this response: " + data.value)
         
 # @spaces.GPU
-def generate_response(message, chatbot, system_prompt="",):
+def generate_response(message, history):
     set_seed(SEED)
-
-    tokenized_prompt = tokenizer(message, return_tensors="pt", padding=True, truncation=True, max_length=128)
-    print(tokenized_prompt)
-
-    output_sequences = model.generate(**tokenized_prompt, max_length=1024, num_return_sequences=1)
-    decoded_output = tokenizer.decode(output_sequences[0], skip_special_tokens=True)
-    print(decoded_output)
-
-    yield decoded_output
-
-chatbot_stream = gr.Chatbot()
+    # data = {
+    #     "role": "user", 
+    #     "content": message
+    # }
+    # tokenized_prompt = tokenizer(data["content"], return_tensors="pt", padding=True, truncation=True, max_length=128)
+    # print(tokenized_prompt)
+
+    # output_sequences = model.generate(**tokenized_prompt, max_length=1024, num_return_sequences=1)
+    # decoded_output = tokenizer.batch_decode(output_sequences, skip_special_tokens=True)[0]
+    # print(decoded_output)
+
+    # yield decoded_output
+    temperature=0.4
+    top_p=0.95
+    top_k=50
+    max_new_tokens=256
+
+    message_template = [
+        {
+            "role": "system",
+            "content": "You are a highly knowledgeable and friendly chatbot equipped with extensive information across various domains. Your goal is to understand and respond to user inquiries with accuracy and clarity. You're adept at providing detailed explanations, concise summaries, and insightful responses. Your interactions are always respectful, helpful, and focused on delivering the most relevant information to the user.",
+        },
+        {"role": "user", "content": message},
+    ]
+
+    # Set tokenize correctly. Otherwise ticking the box breaks it.
+    prompt = pipe.tokenizer.apply_chat_template(message_template, tokenize=False, add_generation_prompt=True)
+    outputs = pipe(prompt, max_new_tokens=max_new_tokens, do_sample=True,
+                   temperature=temperature, top_k=top_k, top_p=top_p, repetition_penalty=1.10)
+
+    print(outputs[0]["generated_text"])
+    return outputs[0]["generated_text"]
+
+# chatbot_stream = gr.Chatbot()
 chat_interface_stream = gr.ChatInterface(generate_response, 
                  title=title, 
                  description=description, 
                  textbox=gr.Textbox(),
-                 chatbot=chatbot_stream,
+                #  chatbot=chatbot_stream,
                  css=css, 
                  examples=examples, 
                  #cache_examples=True, 
@@ -74,4 +98,17 @@ with gr.Blocks() as demo:
     chat_interface_stream.render()
         
 if __name__ == "__main__":
-    demo.queue().launch(share=True)
+    demo.queue().launch(share=True)
+
+
+# messages = [
+#     {
+#         "role": "system",
+#         "content": "You are a friendly chatbot who always responds in the style of a thug",
+#     },
+#     {"role": "user", "content": "How many helicopters can a human eat in one sitting?"},
+# ]
+# model_inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to("cuda")
+# input_length = model_inputs.shape[1]
+# generated_ids = model.generate(model_inputs, do_sample=True, max_new_tokens=20)
+# print(tokenizer.batch_decode(generated_ids[:, input_length:], skip_special_tokens=True)[0])