app.py

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer, GenerationConfig

tokenizer = AutoTokenizer.from_pretrained("ShieldX/manovyadh-1.1B-v1")
model = AutoModelForCausalLM.from_pretrained("ShieldX/manovyadh-1.1B-v1")
model = model.to('cpu')

title = "🌱 ManoVyadh 🌱"
description = "Mental Health Counselling Chatbot"
examples = [["I have been feeling more and more down for over a month. I have started having trouble sleeping due to panic attacks, but they are almost never triggered by something that I know of.", "I self-harm, and I stop for awhile. Then when I see something sad or depressing, I automatically want to self-harm.", "I am feeling sad for my friend's divorce"]]

def predict(message, history):
    def formatted_prompt(question)-> str:
      sysp = "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."
      return f"<|im_start|>system\n{sysp}<|im_end|>\n<|im_start|>user\n{question}<|im_end|>\n<|im_start|>assistant:"

    # history_transformer_format = history + [[formatted_prompt(message), "", message]]

    # messages = "".join(["".join([f"\n user: {item[2]}, \n assistant: {item[1]}"])  #curr_system_message +
    #             for item in history_transformer_format])

    messages = formatted_prompt(message)

    inputs = tokenizer([messages], return_tensors="pt").to("cpu")

    # streamer = TextStreamer(tokenizer)
    
    generation_config = GenerationConfig(
      penalty_alpha=0.6,
      early_stopping=True, 
      num_beams=4,
      do_sample=True,
      top_k=5,
      temperature=0.5,
      repetition_penalty=1.2,
      max_new_tokens=64,
      eos_token_id=tokenizer.eos_token_id,
      pad_token_id=tokenizer.eos_token_id
  )
    
    outputs = model.generate(**inputs, generation_config=generation_config)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response[len(formatted_prompt(message)):]

    # partial_message  = ""
    # for i in response:
    #   if response!="" or i!="":
    #     partial_message+=i
    #     yield partial_message

gr.ChatInterface(
    predict, 
    title=title,
    description=description,
    examples=examples,
    theme="finlaymacklon/boxy_violet",
    ).launch(debug=True)