Upload Tiny-Alice-multi-turn-chat.py

Tiny-Alice-multi-turn-chat.py

@@ -0,0 +1,74 @@
+import os
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# Load model and tokenizer
+model_path = "D1rtyB1rd/Dirty-Alice-Tiny-1.1B-V2"
+model = AutoModelForCausalLM.from_pretrained(model_path)
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+
+# Define the stop token and system message
+stop_token_id = 2 ## </s>
+system_message = "<|system|>\nYou are Alice.\n</s>"
+
+def chat_with_model(prompt_text, stop_token_id, model, tokenizer):
+    # Encode the prompt text
+    encoded_prompt = tokenizer.encode(prompt_text, add_special_tokens=False, return_tensors="pt")
+
+    # Generate response
+    output_sequences = model.generate(
+        input_ids=encoded_prompt,
+        max_new_tokens=1024,
+        temperature=0.2,
+        repetition_penalty=1.2,
+        top_k=20,
+        top_p=0.9,
+        do_sample=True,
+        num_return_sequences=1,
+        eos_token_id=stop_token_id,
+    )
+
+    # Decode the generated sequence
+    generated_sequence = output_sequences[0].tolist()
+    text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)
+
+    # Find the position of the stop token and truncate if necessary
+    stop_token_str = tokenizer.decode([stop_token_id], clean_up_tokenization_spaces=True)
+    if stop_token_str in text:
+        text = text.split(stop_token_str)[0]  # Remove text after the stop token
+
+    response_text = text[len(prompt_text):].strip()  # Extract only the response text
+    return response_text
+
+def build_prompt(conversation_history, user_input):
+    """
+    Constructs the prompt for the model using conversation history and the latest user input.
+    """
+    prompt_text = f"{conversation_history}\n<|user|>\n{user_input}\n</s>\n<|assistant|>\n"
+    return prompt_text
+
+def main():
+    # Initialize conversation history with the system message
+    conversation_history = f"{system_message}\n"
+
+    # Chat loop
+    while True:
+        user_input = input("User: ")  # Get text input from the user
+
+        # Construct prompt text for model input
+        prompt_text = build_prompt(conversation_history, user_input)
+        
+        response_text = chat_with_model(prompt_text, stop_token_id, model, tokenizer)
+        response_text = response_text.replace('<s>', '')
+
+        print(f"\n------\nAlice:\n{response_text}\n------")
+
+        # Update conversation history
+        conversation_history += f"<|user|>\n{user_input}\n</s>\n<|assistant|>\n{response_text}\n</s>\n"
+
+        # Trim the conversation history to avoid overly long inputs
+        if len(conversation_history) > 2048:
+            conversation_history = conversation_history[-1024:]
+
+if __name__ == "__main__":
+    main()
+