Update src/app/main_agent.py

src/app/main_agent.py

@@ -57,15 +57,8 @@ import torch
 from transformers import pipeline
 import os
 
-model_name = 'meta-llama/Llama-3.1-8B' #"google/gemma-3-4b-it"
 # Load the Gemma 3 model pipeline once
-gemma_pipeline = pipeline(
-    task="text-generation",
-    model= model_name,  # or your preferred Gemma 3 model
-    device=0,  # set -1 for CPU, 0 or other for GPU
-    torch_dtype=torch.bfloat16,
-    use_auth_token=os.getenv("HF_TOKEN")
-)
+pipe = pipeline("text-generation", model="google/gemma-3-1b-it", device="cuda", torch_dtype=torch.bfloat16)
 
 def create_agent(accent_tool_obj) -> tuple[Runnable, Runnable]:
     accent_tool = Tool(
@@ -94,15 +87,19 @@ def create_agent(accent_tool_obj) -> tuple[Runnable, Runnable]:
     def follow_up_node(messages: list[BaseMessage]) -> AIMessage:
         user_question = messages[-1].content
         transcript = accent_tool_obj.last_transcript or ""
-        prompt = f"""You are given this transcript of a video:
-
-\"\"\"{transcript}\"\"\"
-
-Now respond to the user's follow-up question: {user_question}
-"""
-        # Use the pipeline to generate the response text
-        # pipeline output is a list of dicts with 'generated_text'
-        outputs = gemma_pipeline(prompt, max_new_tokens=256, do_sample=False)
+        messages = [
+            [
+                {
+                    "role": "system",
+                    "content": [{"type": "text", "text": "You are a helpful assistant."},]
+                },
+                {
+                    "role": "user",
+                    "content": [{"type": "text", "text": "Analyse the transcript. "},]
+                },
+            ],
+        ]
+        outputs = pipe(prompt, max_new_tokens=256, do_sample=False)
         response_text = outputs[0]['generated_text']
 
         return AIMessage(content=response_text)