import gradio as gr
import time

from transformers import AutoTokenizer
import transformers
import torch


global tokenizer, pipeline

model = "shaneperry0101/Health-Llama-3.2-1B"
tokenizer = AutoTokenizer.from_pretrained(model)
pipeline = transformers.pipeline(
    "text-generation",
    model=model,
    torch_dtype=torch.float16,
    device_map="auto",
)

def chat_response(message):
  system_message="Welcome! You're now communicating with an AI model trained to assist with information about general health disease. Feel free to ask about causes, symptoms, medications, and treatment options!"
  prompt = f"[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n{message}[/INST]"
  sequences = pipeline(
      f'[INST] {prompt} [/INST]',
      do_sample=True,
      top_k=10,
      num_return_sequences=1,
      eos_token_id=tokenizer.eos_token_id,
      max_length=500,
  )
  for seq in sequences:
      bot_message = seq['generated_text'].replace(prompt, '').split('[/INST]')[-1]

  return bot_message

with gr.Blocks() as demo:
    chatbot = gr.Chatbot(label="Personal Health Assistant")
    msg = gr.Textbox()
    clear = gr.ClearButton([msg, chatbot])

    def respond(message, chat_history):
        bot_message = chat_response(message)
        chat_history.append((message, bot_message))
        time.sleep(2)
        return "", chat_history

    msg.submit(respond, [msg, chatbot], [msg, chatbot])


if __name__ == "__main__":
    demo.launch()