load in 8 bit

app.py

@@ -1,30 +1,29 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import transformers
 import torch
 
 model = f"tiiuae/falcon-7b"
     
 tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
 
-generator = transformers.pipeline(
-    "text-generation",
-    model=model,
-    # offload_folder='/tmp',
-    tokenizer=tokenizer,
+model = AutoModelForCausalLM.from_pretrained(
+    model,
     torch_dtype=torch.bfloat16,
-    trust_remote_code=True,
     device_map="auto",
+    load_in_8bit=True,
 )
 
-def greet(name):
-    v = generator( 
-            f""" {name}""",
-            max_length=25,
-            do_sample=False,
-            num_return_sequences=1,
-            eos_token_id=tokenizer.eos_token_id,)
-    return v[0]['generated_text']
+def greet(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    v = model.generate(
+        input_ids=inputs["input_ids"],
+        attention_mask=inputs["attention_mask"],
+        do_sample=True,
+        temperature=0.6,
+        top_p=0.9,
+        max_new_tokens=50,
+    )   
+    return tokenizer.decode(v[0].to("cpu"))
 
 iface = gr.Interface(fn=greet, inputs="text", outputs="text")
 iface.launch()