burtenshaw HF staff commited on
Commit
1dfdd1b
·
1 Parent(s): 6222a64

add text only mode

Browse files
Files changed (1) hide show
  1. app/app.py +12 -1
app/app.py CHANGED
@@ -34,6 +34,7 @@ LANGUAGES: dict[str, str] = {
34
 
35
  BASE_MODEL = os.getenv("MODEL", "meta-llama/Llama-3.2-11B-Vision-Instruct")
36
  ZERO_GPU = bool(os.getenv("ZERO_GPU", False)) or True if str(os.getenv("ZERO_GPU")).lower() == "true" else False
 
37
 
38
  def create_inference_client(
39
  model: Optional[str] = None, base_url: Optional[str] = None
@@ -48,7 +49,7 @@ def create_inference_client(
48
  """
49
  if ZERO_GPU:
50
  tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
51
- model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, load_in_4bit=True)
52
  return pipeline("text-generation", model=model, tokenizer=tokenizer)
53
  else:
54
  return InferenceClient(
@@ -91,6 +92,16 @@ def format_history_as_messages(history: list):
91
  current_role = None
92
  current_message_content = []
93
 
 
 
 
 
 
 
 
 
 
 
94
  for entry in history:
95
  content = entry["content"]
96
 
 
34
 
35
  BASE_MODEL = os.getenv("MODEL", "meta-llama/Llama-3.2-11B-Vision-Instruct")
36
  ZERO_GPU = bool(os.getenv("ZERO_GPU", False)) or True if str(os.getenv("ZERO_GPU")).lower() == "true" else False
37
+ TEXT_ONLY = bool(os.getenv("TEXT_ONLY", False)) or True if str(os.getenv("TEXT_ONLY")).lower() == "true" else False
38
 
39
  def create_inference_client(
40
  model: Optional[str] = None, base_url: Optional[str] = None
 
49
  """
50
  if ZERO_GPU:
51
  tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
52
+ model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, load_in_8bit=True)
53
  return pipeline("text-generation", model=model, tokenizer=tokenizer)
54
  else:
55
  return InferenceClient(
 
92
  current_role = None
93
  current_message_content = []
94
 
95
+ if TEXT_ONLY:
96
+ for entry in history:
97
+ messages.append(
98
+ {
99
+ "role": entry["role"],
100
+ "content": entry["content"]
101
+ }
102
+ )
103
+ return messages
104
+
105
  for entry in history:
106
  content = entry["content"]
107