Create drift_agent.py

drift_agent.py

@@ -0,0 +1,17 @@
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
+
+MODEL_NAME = "google/flan-t5-base"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
+model.eval()
+
+def generate_response(prompt: str) -> str:
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=160
+        )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)