Update app.py

app.py

@@ -1,14 +1,21 @@
-from transformers import pipeline,GemmaForCausalLM,AutoTokenizer
+from transformers import pipeline,GemmaForCausalLM,AutoTokenizer,BitsAndBytesConfig
 import gradio as gr
 import spaces
+import torch
 # ignore_mismatched_sizes=True
+quantization_config = BitsAndBytesConfig(load_in_4bit=True)
 tokenizer = AutoTokenizer.from_pretrained('google/gemma-2-9b')
-model = GemmaForCausalLM.from_pretrained('google/gemma-2-9b',ignore_mismatched_sizes=True)
-pipe = pipeline('text-generation', model=model,tokenizer = tokenizer)
+model = GemmaForCausalLM.from_pretrained('google/gemma-2-9b',
+                                        quantization_config=quantization_config
+                                        )
+# pipe = pipeline('text-generation', model=model,tokenizer = tokenizer)
 
 @spaces.GPU(duration=120)
 def generate(prompt):
-    return pipe(prompt)[0]['generated_text']
+    input_ids = tokenizer(prompt, return_tensors="pt").to("cuda")
+    outputs = model.generate(**input_ids)
+    return tokenizer.decode(outputs[0]);
+    # return pipe(prompt)[0]['generated_text']
 
 gr.Interface(
     fn=generate,