import gradio as gr
import os

MY_TOKEN = os.getenv('MY_TOKEN')

DESCRIPTION = "# Meditron-7B v0.2"
from transformers import AutoModelForCausalLM, pipeline, AutoTokenizer

model_id = "epfl-llm/meditron-7b"
model = AutoModelForCausalLM.from_pretrained(model_id, token=MY_TOKEN)
tokenizer = AutoTokenizer.from_pretrained(model_id, token=MY_TOKEN)
pp = pipeline("text-generation", model=model, tokenizer=tokenizer)

with gr.Blocks(css="style.css") as demo:
    gr.Markdown(DESCRIPTION)
    iface = gr.Interface(fn=pp, inputs="text", outputs="text")

if __name__ == "__main__":
    demo.queue(max_size=20).launch()