Update README.md

README.md

@@ -35,25 +35,35 @@ This is the model card of a 🤗 transformers model that has been pushed on the
 
 ## Uses
 
-use under gen function to parse output from the LLM
-Loaded function of LLM is as the same as other LLM
-
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import re
 
+model_id = "jaeyoungk/albatross"  # safetensors 컨버팅된 레포
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
+tokenizer = AutoTokenizer.from_pretrained('meta-llama/Meta-Llama-3-8B-Instruct')
+model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map='auto')
+
+
 def gen(x):
     system_prompt = f"""
     Make a trading decision based on the following data.
     Please respond with a JSON object in the following format: 
     {{"investment_decision": string, "summary_reason": string, "short_memory_index": number, "middle_memory_index": number, "long_memory_index": number, "reflection_memory_index": number}}
     investment_decision must always be one of {{buy, sell, hold}}
-    Print the memory index value to 4 decimal places. If it exceeds, round up.
     """
 
     # Tokenizing the input and generating the output
    
     inputs = tokenizer(
     [
-        f"system{system_prompt}user{x}"
+        f"<|start_header_id|>system<|end_header_id|>{system_prompt}<|eot_id|><|start_header_id|>user<|end_header_id|>{x}<|end_header_id|>"
     ], return_tensors = "pt").to("cuda")
 
     
@@ -66,36 +76,27 @@ def gen(x):
     
     full_text = tokenizer.decode(gened[0])
     
-    # Define possible start phrases
-    possible_start_phrases = ["{\"investment_decision\": \"buy\"", "{\"investment_decision\": \"sell\"", "{\"investment_decision\": \"hold\""]
-    start_idx = -1
-
-    # Find the index for the start phrase
-    for phrase in possible_start_phrases:
-        start_idx = full_text.find(phrase)
-        if start_idx != -1:
-            break
-
-    if start_idx == -1:
-        return "No valid investment decision found in the output."
-
-    # Find the index for the end phrase
-    end_phrase = "\"reflection_memory_index\":"
-    end_idx = full_text.find(end_phrase, start_idx)
-
-    if end_idx == -1:
-        return "No valid reflection_memory_index found in the output."
-
-    # Find the end of the reflection_memory_index value
-    end_idx = full_text.find('}', end_idx)
-    if end_idx == -1:
-        return "No closing bracket found in the output."
+    # Finding the second occurrence of 'user<|end_header_id|'
+    start_phrase = "user<|end_header_id|>"
+    first_occurrence = full_text.find(start_phrase)
+    second_occurrence = full_text.find(start_phrase, first_occurrence + len(start_phrase))
+    
+    if second_occurrence == -1:
+        # If the second occurrence is not found, fallback to using the first occurrence
+        start_idx = first_occurrence + len(start_phrase)
+    else:
+        start_idx = second_occurrence + len(start_phrase)
+    
+    # Find the index of the next special token after the start index
+    end_idx = full_text.find('\\<|eot_id|', start_idx)
 
     # Extract the text between start_idx and end_idx
-    extracted_text = full_text[start_idx:end_idx+1].strip()
+    extracted_text = full_text[start_idx:end_idx].strip()
 
     return extracted_text
 
+gen('input your text here')
+
 
 
 ### Direct Use