add

albatross_classifier_showcase.py

@@ -0,0 +1,61 @@
+import torch
+from torch import cuda
+
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel, PeftConfig
+
+import streamlit as st
+from streamlit_chat import message
+from huggingface_hub import login
+
+login(token='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')
+device = f'cuda:{cuda.current_device()}' if cuda.is_available() else 'cpu'
+
+peft_model_id = "DMPark/albatross_classifier"
+config = PeftConfig.from_pretrained(peft_model_id)
+# model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, return_dict=True, load_in_8bit=True, device_map='auto')
+model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path)
+tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
+
+model = PeftModel.from_pretrained(model, peft_model_id)
+
+def evaluate_model(model, tokenizer, device, instruction, input=None, temperature=0.1, top_p=0.75, top_k=40, num_beams=4, max_new_tokens=128):
+    prompt_instruction = "### Instruction:\n{}\n" + (f"### Input:\n{input}\n" if input else "") + "### Response:\n"
+    prompt = "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n" + prompt_instruction.format(instruction)
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    with torch.no_grad():
+        generation_output = model.generate(
+            input_ids=inputs["input_ids"],
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            num_beams=num_beams,
+            max_length=max_new_tokens + inputs["input_ids"].shape[1]
+        )
+    
+    output = tokenizer.decode(generation_output[0])
+    return output.split("### Response:")[1].strip()
+
+st.header("🦜Albatross")
+
+if 'generated' not in st.session_state:
+    st.session_state['generated'] = []
+
+if 'past' not in st.session_state:
+    st.session_state['past'] = []
+
+    
+with st.form('form', clear_on_submit=True):
+    user_input = st.text_area('You: ', '')
+    submitted = st.form_submit_button('Send')
+
+if submitted and user_input:
+    output = evaluate_model(model, tokenizer, device, instruction="Predict FOMC' stance toward Interest Rate Dicision (Hawkish/Neutral/Dovish)", input=fomc)
+    st.session_state.past.append(user_input)
+    st.session_state.generated.append(output)
+
+if 'generated' in st.session_state:
+    for i in range(len(st.session_state['generated'])-1, -1, -1):
+        message(st.session_state['past'][i], is_user=True, key=str(i) + '_user')
+        message(st.session_state["generated"][i], key=str(i))

requirements.txt

@@ -0,0 +1,6 @@
+torch
+transformers
+peft
+huggingface_hub
+streamlit-chat
+streamlit