tairov / llama2.mojo

Inference Llama 2 in one file of pure 🔥
https://www.modular.com/blog/community-spotlight-how-i-built-llama2-by-aydyn-tairov
MIT License
2.09k stars 140 forks source link

Prompt input? #11

Closed radames closed 1 year ago

radames commented 1 year ago

Hi @tairov great work here! I notice you're not using the prompt, it would be fun to be able to input it, I could add it to the #10

tairov commented 1 year ago

Hi @radames , thanks for sending PR, I'll take a look tomorrow. Yeah, I didn't have a chance to implement one more fn for prompt support. Would you mind to create another PR for prompt input? thanks!

radames commented 1 year ago

If you implement the step on Mojo, I can also add the Gradio UI part of it! I couldn't find how to get the args using argv 🤦🏼 But this worked on my tests

    let os = Python.import_module("os")
    let prompt = os.environ.get('PROMPT', "")
tairov commented 1 year ago

Probably that's why I haven't ported prompt part right away from llama2.py , because it was not clear how to read args I would implement it like in llama2.c with arg names .

trevorhobenshield commented 1 year ago

You can read in args like this:

from sys import argv

fn main() raises:
    let args = argv()
    print(args[1])
mojo main.mojo foo
tairov commented 1 year ago

prompt support added @radames could you add changes to gradio app as well? Is it possible to execute it cli in this way via gradio app?

mojo llama2.mojo stories15M.bin -s 100 -n 256 -t 0.5 -i "Llama is an animal"
radames commented 1 year ago

awesome @tairov I'll do it as soon as I can!