Update README.md

README.md

@@ -47,6 +47,7 @@ text = """<schema>{schema}</schema>
 <question>{question}</question>
 <sql>"""
 ```
+pytorch
 
 ```python
 from transformers import AutoModelForCasualLM, AutoTokenizer
@@ -58,6 +59,14 @@ inputs = tokenizer(text, return_tensors="pt")
 outputs = model.generate(**inputs, max_new_tokens=200)
 print(tokenizer.decode(outputs[0], skip_special_tokens=True).split('<sql>')[1].split('</sql>')[0])
 ```
+flax
+
+```python
+from transformers import AutoModelForCasualLM, AutoTokenizer
+device = "cuda"
+model = FlaxAutoModelForCausalLM.from_pretrained("PipableAI/pipSQL1b" , from_pt=True)
+tokenizer = AutoTokenizer.from_pretrained("PipableAI/pipSQL1b")
+```
 
 ## The PipableAI team