Update app.py

app.py

@@ -1,36 +1,34 @@
 import gradio as gr
 
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import (
+    AutoModelForCausalLM,
+    AutoConfig,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+)
 import transformers
 import torch
 
-model = "tiiuae/falcon-40b"
+model_name = "tiiuae/falcon-40b"
 
-tokenizer = AutoTokenizer.from_pretrained(model)
-pipeline = transformers.pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    torch_dtype=torch.bfloat16,
+config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model,
+    quantization_config=BitsAndBytesConfig(load_in_4bit=True),
     trust_remote_code=True,
-    load_in_4bit=True,
+    torch_dtype=torch.bfloat16,
     device_map="auto",
 )
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-def falcon(input_text):
 
-    sequences = pipeline(
-        input_text, # "Was ist das höchste Gebäude in der Welt?"
-        max_length=200,
-        do_sample=True,
-        top_k=10,
-        num_return_sequences=1,
-        eos_token_id=tokenizer.eos_token_id,
-    )
-    for seq in sequences:
-        print(f"Result: {seq['generated_text']}")
+def falcon(input_text):
+    input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
+    outputs = model.generate(input_ids, max_length=100, do_sample=True, top_k=10)
+    decoded = tokenizer.decode(outputs[0])
+    
+    return decoded
 
-    return sequences[0]['generated_text']
 
 iface = gr.Interface(fn=falcon, inputs="text", outputs="text")
-iface.launch() # To create a public link, set `share=True`
+iface.launch()  # To create a public link, set `share=True`