changed to ft model

app.py

@@ -6,16 +6,16 @@ import threading
 import streamlit as st
 import random
 from typing import Iterable
-# from unsloth import FastLanguageModel
+from unsloth import FastLanguageModel
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, PreTrainedTokenizerFast
 from datetime import datetime
 from threading import Thread
 
-# fine_tuned_model_name = "jed-tiotuico/twitter-llama"
-# sota_model_name = "unsloth/mistral-7b-instruct-v0.2-bnb-4bit"
+fine_tuned_model_name = "jed-tiotuico/twitter-llama"
+sota_model_name = "unsloth/mistral-7b-instruct-v0.2-bnb-4bit"
+# fine_tuned_model_name = "MBZUAI/LaMini-GPT-124M"
+# sota_model_name = "MBZUAI/LaMini-GPT-124M"
 
-fine_tuned_model_name = "MBZUAI/LaMini-GPT-124M"
-sota_model_name = "MBZUAI/LaMini-GPT-124M"
 alpaca_input_text_format = "### Instruction:\n{}\n\n### Response:\n"
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 # if device is cpu try mps?
@@ -24,16 +24,15 @@ if device == "cpu":
     device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
 
 def get_model_tokenizer(sota_model_name):
-    tokenizer = AutoTokenizer.from_pretrained(
-        sota_model_name,
-        cache_dir="/data/.hf_cache",
-        trust_remote_code=True
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name = "jed-tiotuico/twitter-llama",
+        max_seq_length = 200,
+        dtype = None,
+        load_in_4bit = True,
+        cache_dir = "/data/.cache/hf-models",
+        token=st.secrets["HF_TOKEN"]
     )
-    model = AutoModelForCausalLM.from_pretrained(
-        sota_model_name,
-        cache_dir="/data/.hf_cache",
-        trust_remote_code=True
-    ).to(device)
+    FastLanguageModel.for_inference(model)
 
     return model, tokenizer
 
@@ -61,16 +60,14 @@ def write_stream_user_chat_message(user_chat, model, token, prompt):
         return new_customer_msg
 
 def get_mistral_model_tokenizer(sota_model_name):
-    tokenizer = AutoTokenizer.from_pretrained(
-        sota_model_name,
-        cache_dir="/data/.hf_cache",
-        trust_remote_code=True
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name = "unsloth/mistral-7b-instruct-v0.2-bnb-4bit",
+        max_seq_length = max_seq_length,
+        dtype = dtype,
+        load_in_4bit = load_in_4bit,
+        cache_dir = "/data/.cache/hf-models",
     )
-    model = AutoModelForCausalLM.from_pretrained(
-        sota_model_name,
-        cache_dir="/data/.hf_cache",
-        trust_remote_code=True
-    ).to(device)
+    FastLanguageModel.for_inference(model)
 
     return model, tokenizer