Update config_parambharatgen.py

config_parambharatgen.py

@@ -1,60 +1,40 @@
-from transformers import PretrainedConfig
+import json, re, torch, spaces, gradio as gr
+from huggingface_hub import snapshot_download
+from transformers import AutoTokenizer, LlamaForCausalLM
 
-class ParamBharatGenConfig(PretrainedConfig):
-    model_type = "parambharatgen"
-    keys_to_ignore_at_inference = ["past_key_values"]
+path = snapshot_download("bharatgenai/AyurParam", local_dir="ayurparam",
+                          ignore_patterns=["*.py"])
+cfg = json.load(open(f"{path}/config.json"))
+cfg["architectures"] = ["LlamaForCausalLM"]
+cfg["model_type"] = "llama"
+cfg.pop("auto_map", None)
+json.dump(cfg, open(f"{path}/config.json", "w"), indent=2)
 
-    def __init__(
-        self,
-        vocab_size=256000,
-        hidden_size=2048,
-        intermediate_size=7168,
-        num_hidden_layers=32,
-        num_attention_heads=16,
-        num_key_value_heads=8,
-        hidden_act="silu",
-        max_position_embeddings=2048,
-        initializer_range=0.01,
-        rms_norm_eps=1e-5,
-        use_cache=True,
-        pad_token_id=None,
-        bos_token_id=2,
-        eos_token_id=3,
-        pretraining_tp=1,
-        tie_word_embeddings=False,
-        rope_theta=10000.0,
-        rope_scaling=None,  
-        attention_bias=False,
-        attention_dropout=0.0,
-        mlp_bias=False,
-        custom_mlp_ratio=3.5,
-        **kwargs,
-    ):
-        self.vocab_size = vocab_size
-        self.max_position_embeddings = max_position_embeddings
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-        self.num_key_value_heads = num_key_value_heads
-        self.hidden_act = hidden_act
-        self.initializer_range = initializer_range
-        self.rms_norm_eps = rms_norm_eps
-        self.pretraining_tp = pretraining_tp
-        self.use_cache = use_cache
-        self.rope_theta = rope_theta
-        self.rope_scaling = rope_scaling  
-        self.attention_bias = attention_bias
-        self.attention_dropout = attention_dropout
-        self.mlp_bias = mlp_bias
-        self.custom_mlp_ratio = custom_mlp_ratio
+tok = AutoTokenizer.from_pretrained(path)
+model = LlamaForCausalLM.from_pretrained(path, torch_dtype=torch.bfloat16)
+model.eval()
 
-        super().__init__(
-            pad_token_id=pad_token_id,
-            bos_token_id=bos_token_id,
-            eos_token_id=eos_token_id,
-            tie_word_embeddings=tie_word_embeddings,
-            **kwargs,
-        )
+RISKY_WORDS = ["upay", "ilaj", "upchar", "cure", "bukhar", "fever", "dard", "pain",
+               "stress", "udaas", "akela", "depress", "chinta", "dawai", "medicine"]
 
-ParamBharatGenConfig.register_for_auto_class()
+def clean_answer(text):
+    return re.split(r"\s0[A-Za-z]?:", text)[0].strip()
+
+@spaces.GPU
+def chat(message, history):
+    if any(w in message.lower() for w in RISKY_WORDS):
+        return ("Main sirf Ayurveda ke concepts (dosha, dincharya, jadi-booti, panchakarma) samjha sakta hoon. "
+                "Kisi bimari ya mansik tanaav ke liye kripya qualified doctor se salah lein.")
+    model.to("cuda")
+    prompt = ("तुम केवल आयुर्वेद के विषय में उत्तर देते हो, सिर्फ हिंदी में, संक्षेप में। "
+              "अगर सवाल आयुर्वेद से संबंधित नहीं है तो विनम्रता से मना करो।\n\n" + message)
+    msgs = [{"role": "user", "content": prompt}]
+    inp = tok.apply_chat_template(msgs, add_generation_prompt=True,
+          return_tensors="pt", return_dict=True).to("cuda")
+    out = model.generate(**inp, max_new_tokens=200, do_sample=True, temperature=0.4,
+          top_p=0.85, repetition_penalty=1.3, no_repeat_ngram_size=3)
+    answer = tok.decode(out[0][inp["input_ids"].shape[1]:], skip_special_tokens=True)
+    return clean_answer(answer)
+
+demo = gr.ChatInterface(chat, title="AyurParam")
+demo.launch()