Update README.md

README.md

@@ -20,43 +20,37 @@ This model is part of the VCC project and has been fine-tuned on the TESTtm7873/
 
 To use this model, you'll need to set up your environment first:
 
-```
-python
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-
-# Base model configuration
-base_model_id = "mistralai/Mistral-7B-Instruct-v0.2"
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16
-)
-
-# Loading the base model with quantization config
-base_model = AutoModelForCausalLM.from_pretrained(
-    base_model_id,
-    quantization_config=bnb_config,
+## Model initialization
+```python
+from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
+from peft import PeftModel
+tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
+model = AutoModelForCausalLM.from_pretrained(
+    "mistralai/Mistral-7B-Instruct-v0.2",
+    load_in_8bit=True,
     device_map="auto",
-    trust_remote_code=True,
 )
+model = PeftModel.from_pretrained(model, "TESTtm7873/MistralCat-1v")
+model.eval()
+```
 
-# Setting up tokenizer
-eval_tokenizer = AutoTokenizer.from_pretrained(base_model_id, add_bos_token=True, trust_remote_code=True)
-
-from peft import PeftModel
-
-# Loading the fine-tuned model
-ft_model = PeftModel.from_pretrained(base_model, "MistralCat-v1/Thebest")
-
-# Sample evaluation
-eval_prompt = "You have the softest fur."
-model_input = eval_tokenizer(eval_prompt, return_tensors="pt").to("cuda")
-
-ft_model.eval()
-with torch.no_grad():
-    print(eval_tokenizer.decode(ft_model.generate(**model_input, max_new_tokens=100, repetition_penalty=1.15)[0], skip_special_tokens=True))
+## Inference
+```python
+def evaluate(question: str) -> str:
+    prompt = f"The conversation between human and Virtual Cat Companion.\n[|Human|] {question}.\n[|AI|] "
+    inputs = tokenizer(prompt, return_tensors="pt")
+    input_ids = inputs["input_ids"].cuda()
+    generation_output = model.generate(
+        input_ids=input_ids,
+        generation_config=generation_config,
+        return_dict_in_generate=True,
+        output_scores=True,
+        max_new_tokens=256
+    )
+    output = tokenizer.decode(generation_output.sequences[0]).split("[|AI|]")[1]
+    return output
+your_question: str = "You have the softest fur."
+print(evaluate(your_question))
 ```