Update app.py

app.py

@@ -1,10 +1,53 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+# """
+# For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
+# """
+# client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+
+from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM#, MambaForCausalLM
+from peft import PeftConfig, PeftModel
+
+config = PeftConfig.from_pretrained("jonathanjordan21/mos-mamba-6x130m-trainer")
+
+tokenizer = AutoTokenizer.from_pretrained("jonathanjordan21/mos-mamba-6x130m-trainer", trust_remote_code=True)
+
+model = AutoModelForCausalLM.from_pretrained(
+    "jonathanjordan21/mos-mamba-6x130m-trainer",
+    eos_token_id=tokenizer.eos_token_id,
+    trust_remote_code=True
+)
+
+model = PeftModel.from_pretrained(model, "jonathanjordan21/mos-mamba-6x130m-trainer",)#, adapter_name="norobots")
+model = model.merge_and_unload()
+
+
+def invoke(
+    message,
+    history: list[tuple[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+):
+    messages = [{"role": "system", "content": system_message}]
+
+    for val in history:
+        if val[0]:
+            messages.append({"role": "user", "content": val[0]})
+        if val[1]:
+            messages.append({"role": "assistant", "content": val[1]})
+
+    messages.append({"role": "user", "content": message})
+
+    tokens = tokenizer.apply_chat_template(message, return_tensors='pt')
+
+    out = model.generate(**tokens, eos_token_id=model.config.eos_token_id, max_new_tokens=max_tokens)
+
+    res = tokenizer.batch_decode(out)
+
+    return res
 
 
 def respond(