Testing falcon for chat completion

postgresml · santiatpml · Aug 8, 2023 · Aug 7, 2023 · Aug 7, 2023 · Aug 7, 2023
commit 18537ab93f0eff20b18a377f5e43a019d73359ae
diff --git a/pgml-apps/pgml-chat/pgml_chat/falcon.py b/pgml-apps/pgml-chat/pgml_chat/falcon.py
@@ -0,0 +1,36 @@
+from pgml import Database
+import os
+from dotenv import load_dotenv
+import asyncio
+from rich.pretty import pprint
+# Load .env file
+load_dotenv(".env")
+
+database_url = os.environ.get("DATABASE_URL")
+db = Database(database_url)
+
+
+async def main():
+    prompt = (
+        "Girafatron is obsessed with giraffes, the most glorious animal on the face of this Earth. Giraftron believes all other animals are irrelevant when compared to the glorious majesty of the giraffe.\nDaniel: Hello, Girafatron!\nGirafatron:",
+    )
+    results = await db.transform(
+        task={
+            "task": "text-generation",
+            "model": "tiiuae/falcon-7b-instruct",
+            "trust_remote_code": True,
+            "torch_dtype": "bfloat16",
+            "device_map": "auto",
+        },
+        inputs=prompt,
+        args={
+            "temperature": 0.7,
+            "top_p": 0.9,
+            "max_length" : 256
+        },
+    )
+    pprint(results)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())