Making app CPU compatible.

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+
 @st.cache(allow_output_mutation=True)
 def get_model():
     # load base LLM model and tokenizer
@@ -12,7 +13,7 @@ def get_model():
     model_id,
     low_cpu_mem_usage=True,
     torch_dtype=torch.float16,
-    load_in_4bit=True,
+    # load_in_4bit=True,
     )
 
     return tokenizer, model
@@ -33,10 +34,13 @@ if user_input and button:
 
     ### Response:
     """
-    input_ids = tokenizer(prompt, return_tensors="pt", truncation=True).input_ids
-    outputs = model.generate(input_ids=input_ids, pad_token_id=tokenizer.eos_token_id, max_new_tokens=500, do_sample=True, top_p=0.75, temperature=0.95, top_k=15)
     st.write("Prompt: ", user_input)
+    input = tokenizer([prompt], padding=True, truncation=True, return_tensors="pt")
+    output = model(**input)
+    # input_ids = tokenizer(prompt, return_tensors="pt", truncation=True)
+    # outputs = model.generate(input_ids=input_ids, pad_token_id=tokenizer.eos_token_id, max_new_tokens=500, do_sample=True, top_p=0.75, temperature=0.95, top_k=15)
+
     st.write("**************")
-    st.write(outputs)
+    st.write(output)
 
 

requirements.txt

@@ -1,5 +1,8 @@
 streamlit
 torch
-transformers
-accelerate
-bitsandbytes
+transformers==4.34.0
+accelerate==0.23.0
+bitsandbytes==0.41.1
+trl==0.4.7
+safetensors>=0.3.1
+peft==0.5.0