Update README.md

README.md

@@ -14,27 +14,48 @@ Merge of two SOLAR models. This is an experiment to improve models ability to le
 
 ## 🌅 Code Example
 
+## 🌅 Code Example
+
+Example also available in [colab](https://colab.research.google.com/drive/10FWCLODU_EFclVOFOlxNYMmSiLilGMBZ?usp=sharing)
+
 ```python
-import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-tokenizer = AutoTokenizer.from_pretrained("macadeliccc/SOLAR-math-2x10.7b",load_in_4bit=True)
-model = AutoModelForCausalLM.from_pretrained(
-    "macadeliccc/SOLAR-math-2x10.7b",
-    device_map="auto",
-    torch_dtype=torch.float16,
-)
+def generate_response(prompt):
+    """
+    Generate a response from the model based on the input prompt.
+
+    Args:
+    prompt (str): Prompt for the model.
+
+    Returns:
+    str: The generated response from the model.
+    """
+    # Tokenize the input prompt
+    inputs = tokenizer(prompt, return_tensors="pt")
+    
+    # Generate output tokens
+    outputs = model.generate(**inputs, max_new_tokens=512, eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.pad_token_id)
 
-conversation = [ {'role': 'user', 'content': 'A rectangle has a length that is twice its width and its area is 50 square meters. Find the dimensions of the rectangle.'} ] 
+    # Decode the generated tokens to a string
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-prompt = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
+    return response
 
-inputs = tokenizer(prompt, return_tensors="pt").to(model.device) 
-outputs = model.generate(**inputs, use_cache=True, max_length=4096)
-output_text = tokenizer.decode(outputs[0]) 
-print(output_text)
+
+# Load the model and tokenizer
+model_id = "macadeliccc/SOLAR-math-2x10.7b"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, load_in_4bit=True)
+
+prompt = "Explain the proof of Fermat's Last Theorem and its implications in number theory."
+
+
+print("Response:")
+print(generate_response(prompt), "\n")
 ```
 
+
 ## Evaluations 
 
 TODO