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 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,streamer=streamer, 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)