# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig

tokenizer = AutoTokenizer.from_pretrained("ShieldX/manovyadh-1.1B-v1-chat")
model = AutoModelForCausalLM.from_pretrained("ShieldX/manovyadh-1.1B-v1-chat").to("cuda")
config = AutoConfig.from_pretrained("ShieldX/manovyadh-1.1B-v1-chat")

def format_prompt(q):
  return f"""###SYSTEM: You are an AI assistant that helps people cope with stress and improve their mental health. User will tell you about their feelings and challenges. Your task is to listen empathetically and offer helpful suggestions. While responding, think about the user’s needs and goals and show compassion and support
      ###USER: {q}
      ###ASSISTANT:"""

prompt = format_prompt("I've never been able to talk with my parents. My parents are in their sixties while I am a teenager. I love both of them but not their personalities. I feel that they do not take me seriously whenever I talk about a serious event in my life. If my dad doesn’t believe me, then my mom goes along with my dad and acts like she doesn’t believe me either. I’m a pansexual, but I can’t trust my own parents. I've fought depression and won; however, stress and anxiety are killing me. I feel that my friends don't listen to me. I know they have their own problems, which I do my best to help with. But they don't always try to help me with mine, when I really need them. I feel as if my childhood has been taken from me. I feel as if I have no one whom I can trust.")

import torch
from transformers import GenerationConfig, TextStreamer
from time import perf_counter

# Check for GPU availability
if torch.cuda.is_available():
    device = "cuda"
else:
    device = "cpu"

# Move model and inputs to the GPU (if available)
model.to(device)

inputs = tokenizer(prompt, return_tensors="pt").to(device)

streamer = TextStreamer(tokenizer)

generation_config = GenerationConfig(
    penalty_alpha=0.6,
    do_sample=True,
    top_k=5,
    temperature=0.5,
    repetition_penalty=1.2,
    max_new_tokens=256,
    streamer=streamer,
    pad_token_id=tokenizer.eos_token_id
)

start_time = perf_counter()
outputs = model.generate(**inputs, generation_config=generation_config)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
output_time = perf_counter() - start_time
print(f"Time taken for inference: {round(output_time, 2)} seconds")

@misc{ShieldX/manovyadh-1.1B-v1-chat,
      url={[https://huggingface.co/ShieldX/manovyadh-1.1B-v1-chat](https://huggingface.co/ShieldX/manovyadh-1.1B-v1-chat)},
      title={ManoVyadh},
      author={Rohan Shaw},
      year={2024}, month={Jan}
}