Update app.py

app.py

@@ -2,9 +2,9 @@ import spaces
 import gradio as gr
 from transformers import pipeline
 
-model = "daniel-dona/gemma-3-270m-it"
+model_name = "daniel-dona/gemma-3-270m-it"
 
-pipe = pipeline("text-generation", model=model, device="cuda")
+#pipe = pipeline("text-generation", model=model, device="cuda")
 
 @spaces.GPU
 def respond(
@@ -28,7 +28,7 @@ def respond(
 
     messages.append({"role": "user", "content": message})
 
-    response = pipe(
+    """response = pipe(
         messages,
         max_new_tokens=max_tokens,
         temperature=temperature,
@@ -38,7 +38,43 @@ def respond(
 
     generated_text = response[0]['generated_text']
 
-    yield generated_text
+    yield generated_text"""
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype="auto",
+        device_map="auto"
+    )
+
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        enable_thinking=model_thinking
+    )
+
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+    sample = True
+
+    if model_temperature == 0:
+        sample = False
+
+
+    # conduct text completion
+    generated_ids = model.generate(
+        **model_inputs,
+        max_new_tokens=max_tokens,
+        do_sample=sample,
+        top_p=top_p,
+        temperature=model_temperature
+    )
+    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist() 
+
+    content = tokenizer.decode(output_ids, skip_special_tokens=True).strip("\n")
+
+    return content
 
 
 """