import gradio as gr
from transformers import pipeline

pipe = pipeline(
    "text-generation",
    model="BioMistral/BioMistral-7B"
)

def generate(prompt):

    final_prompt = f"""
You are a professional Evidence-Based Medicine assistant.

Answer the medical question clearly.

Question:
{prompt}

Answer:
"""

    output = pipe(
        final_prompt,
        max_new_tokens=200,
        do_sample=False,
        temperature=0.1
    )

    text = output[0]["generated_text"]

    return text.replace(final_prompt, "")

demo = gr.Interface(
    fn=generate,
    inputs="text",
    outputs="text",
    title="EBM Assistant"
)

demo.launch()