search

google-ai-edge / ai-edge-torch

Supporting PyTorch models with the Google AI Edge TFLite runtime.
Apache License 2.0
228 stars 26 forks source link

fix the tiny llama conversion issue #7

Closed freedomtan closed 2 months ago

freedomtan commented 2 months ago

Fix the issue #6 as suggested by @kaushiksiva07,

BUG=fixes #6

freedomtan commented 2 months ago

FYR. Actually, it's not a Tiny Llama specific issue. I converted Gemma 2B without running into problems. But what I run the C++ text generation example in https://github.com/google-ai-edge/ai-edge-torch/tree/main/ai_edge_torch/generative/examples/c%2B%2B

bazel run -c opt //ai_edge_torch/generative/examples/c++:text_generator_main -- --tflite_model=PATH/gemma_it.tflite  --sentencepiece_model=PATH/tokenizer.model --start_token="<bos>" --stop_token="<eos>" --num_threads=16 --prompt="Write an email:"

I got meaning results. After applying this patch and regenerating the tflite, I got expected results.

advaitjain commented 2 months ago

I pushed a commit updating the formatting. will re-approve and then the PR should be ready to go.