update

README.md

@@ -7,22 +7,58 @@ tags:
 - trl
 - sft
 licence: license
+license: mit
+language:
+- vi
+pipeline_tag: text-generation
 ---
 
-# Model Card for output
+# Introduction
 
-This model is a fine-tuned version of [google/gemma-3-4b-it](https://huggingface.co/google/gemma-3-4b-it).
+This model is a fine-tuned version of [google/gemma-3-4b-it](https://huggingface.co/google/gemma-3-4b-it) optimized for query rewriting based on conversation history to enhance conversational retrieval and question answering.
 It has been trained using [TRL](https://github.com/huggingface/trl).
 
 ## Quick start
 
 ```python
-from transformers import pipeline
+from transformers import AutoTokenizer, Gemma3ForConditionalGeneration
 
-question = "If you had a time machine, but could only go to the past or the future once and never return, which would you choose and why?"
-generator = pipeline("text-generation", model="r1ck/output", device="cuda")
-output = generator([{"role": "user", "content": question}], max_new_tokens=128, return_full_text=False)[0]
-print(output["generated_text"])
+
+MODEL_PATH = "r1ck/gemma-3-4b-it-rw"
+
+model = Gemma3ForConditionalGeneration.from_pretrained(
+    MODEL_PATH, device_map="auto", attn_implementation='eager'
+).eval()
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+
+prompt_template = """Given a conversation between (user, assistant) and a follow-up message from user, your task is to rewrite the follow-up message to a standalone message that captures all relevant context from the conversation.
+Consume the entire conversation and follow-up message and think deeply about it. If the follow-up message is already clear, don't need to rewrite it, just return the original message. The rewritten message must be in Vietnamese.
+
+# Conversation:
+user: Hà Nội có những đặc điểm văn hóa nào nổi bật từ lịch sử?
+assistant: Hà Nội, với lịch sử là kinh đô của Việt Nam, đã hội tụ nhiều tinh hoa văn hóa từ miền Bắc và cả nước. Thành phố này là nơi quy tụ của những nhân vật ưu tú, thương nhân, nghệ nhân, và thợ thủ công lành nghề từ khắp nơi. Họ mang theo phong tục, tập quán địa phương của mình, từ đó tạo nên nét văn hóa đặc trưng cho Hà Nội.
+
+# Follow-up message:
+Điều này mang lại lợi ích gì về du lịch?
+
+# Rewritten message:
+"""
+
+inputs = tokenizer(
+    [prompt_template],
+    return_tensors="pt"
+).to("cuda")
+
+outputs = model.generate(
+    input_ids=inputs.input_ids,
+    attention_mask=inputs.attention_mask,
+    max_new_tokens=2048,
+    eos_token_id=tokenizer.eos_token_id,
+    use_cache=True,
+)
+response = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+print(response[0])
 ```
 
 ## Training procedure