demo chat with ollama

src/main.py

@@ -1,6 +1,25 @@
 def main():
-    print("Hello from tr-steering-allen!")
 
+    import argparse
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--model","-m", default="glm-4.7-flash:latest", help="model used for generation")
+
+    args = parser.parse_args()
+
+    print(args)
+    
+    from ollama import chat
+
+    prompt = "Hi, how are you today?"
+
+    reply = []
+    for _ in range(10):
+        response = chat(model=args.model, messages=[{"role":"user","content":prompt}])
+        reply.append(response.message.content)
+
+    print("\n".join(reply))
 
 if __name__ == "__main__":
     main()