Upload 2 files

app.py

@@ -2,21 +2,35 @@ import spaces
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, TorchAoConfig
 from threading import Thread
-import torch
-from torchao.quantization import Int8DynamicActivationInt8WeightConfig
-import subprocess
+import os, subprocess, torch
+from torchao.quantization import Int4WeightOnlyConfig, Int8DynamicActivationInt8WeightConfig, Float8DynamicActivationFloat8WeightConfig
+from torchao.dtypes import Int4CPULayout
 
-subprocess.run("pip list", shell=True)
 
-quant_config = Int8DynamicActivationInt8WeightConfig()
+#subprocess.run("pip list", shell=True)
+
+IS_COMPILE = False if torch.cuda.is_available() else True
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+# https://huggingface.co/docs/transformers/en/quantization/torchao?examples-CPU=int8-dynamic-and-weight-only
+if torch.cuda.is_available():
+    quant_config = Float8DynamicActivationFloat8WeightConfig()
+else:
+    #quant_config = Int4WeightOnlyConfig(group_size=128, layout=Int4CPULayout())
+    quant_config = Int8DynamicActivationInt8WeightConfig()
+
 quantization_config = TorchAoConfig(quant_type=quant_config)
 #checkpoint = "HuggingFaceTB/SmolLM2-135M-Instruct"
 checkpoint = "unsloth/gemma-3-4b-it"
-device = "cuda" if torch.cuda.is_available() else "cpu"
 tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-model = AutoModelForCausalLM.from_pretrained(checkpoint, torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-                                             device_map=device, quantization_config=quantization_config)
 #model = AutoModelForCausalLM.from_pretrained(checkpoint, torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32).to(device)
+model = AutoModelForCausalLM.from_pretrained(checkpoint, torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+                                             device_map=device, quantization_config=quantization_config).eval()
+if IS_COMPILE:
+    model.generation_config.cache_implementation = "static"
+    input_text = "Warming up."
+    input_ids = tokenizer(input_text, return_tensors="pt").to(device)
+    output = model.generate(**input_ids, max_new_tokens=10, cache_implementation="static")
 
 def get_duration(message, history, system_message, max_tokens, temperature, top_p, duration):
     return duration
@@ -32,7 +46,7 @@ def respond_stream(message, history, system_message, max_tokens, temperature, to
         add_generation_prompt=True,
         return_tensors="pt",
         return_dict=True,
-    ).to(device)
+    ).to(model.device)
 
     streamer = TextIteratorStreamer(
         tokenizer, skip_prompt=True, skip_special_tokens=True
@@ -46,8 +60,10 @@ def respond_stream(message, history, system_message, max_tokens, temperature, to
         temperature=temperature,
         top_p=top_p,
         eos_token_id=tokenizer.eos_token_id,
-        cache_implementation="static",
+        num_beams=1,
+        output_scores=False,
     )
+    if IS_COMPILE: gen_kwargs["cache_implementation"] = "static"
     thread = Thread(target=model.generate, kwargs=gen_kwargs)
     thread.start()
 
@@ -65,12 +81,11 @@ def respond(message, history, system_message, max_tokens, temperature, top_p, du
         messages,
         tokenize=True,
         add_generation_prompt=True,
-        padding=True,
         return_tensors="pt",
         return_dict=True,
-    ).to(device)
+    ).to(model.device)
 
-    outputs = model.generate(
+    gen_kwargs = dict(
         input_ids=inputs["input_ids"],
         #attention_mask=inputs["attention_mask"],
         max_new_tokens=max_tokens,
@@ -78,8 +93,11 @@ def respond(message, history, system_message, max_tokens, temperature, top_p, du
         temperature=temperature,
         top_p=top_p,
         eos_token_id=tokenizer.eos_token_id,
-        cache_implementation="static",
+        num_beams=1,
+        output_scores=False,
     )
+    if IS_COMPILE: gen_kwargs["cache_implementation"] = "static"
+    outputs = model.generate(**gen_kwargs)
 
     gen_ids = outputs[0][inputs["input_ids"].shape[-1]:]
     return tokenizer.decode(gen_ids, skip_special_tokens=True)

requirements.txt

@@ -1,6 +1,7 @@
 huggingface_hub[hf_xet]
 torch
 torchao
+triton
 transformers
 accelerate
 peft