Update src/app/main_agent.py

src/app/main_agent.py

@@ -56,11 +56,11 @@ import re
 import torch
 from transformers import pipeline
 import os
+# Load model directly
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-model_id = "google/gemma-3-4b-it"
-
-# Load the Gemma 3 model pipeline once
-pipe = pipeline("text-generation", model=model_id, use_auth_token=os.getenv("HF_TOKEN"))
+tokenizer = AutoTokenizer.from_pretrained("google/gemma-3-1b-it")
+model = AutoModelForCausalLM.from_pretrained("google/gemma-3-1b-it")
 
 def create_agent(accent_tool_obj) -> tuple[Runnable, Runnable]:
     accent_tool = Tool(
@@ -101,7 +101,15 @@ def create_agent(accent_tool_obj) -> tuple[Runnable, Runnable]:
                 },
             ],
         ]
-        outputs = pipe(prompt, max_new_tokens=256, do_sample=False)
+        inputs = tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt",
+        )
+        outputs =  model.generate(**inputs, max_new_tokens=64)
+        outputs = tokenizer.batch_decode(outputs)
         response_text = outputs[0]['generated_text']
 
         return AIMessage(content=response_text)