Update app.py

app.py

@@ -1,10 +1,39 @@
 from huggingface_hub import InferenceClient
 import gradio as gr
+import torch
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import transformers
 
-client = InferenceClient(
-    "pyakhurel/mistral-7b-mj-finetuned"
+adapters_name = "pyakhurel/mistral-7b-mj-finetuned"
+model_name = "bn22/Mistral-7B-Instruct-v0.1-sharded"
+device = "cuda"
+
+bnb_config = transformers.BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    load_in_4bit=True,
+    torch_dtype=torch.bfloat16,
+    quantization_config=bnb_config,
+    device_map='auto'
 )
 
+model = PeftModel.from_pretrained(model, adapters_name)
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+tokenizer.bos_token_id = 1
+
+stop_token_ids = [0]
+
+
+
+
 
 def format_prompt(message, history):
   prompt = "<s>"
@@ -33,7 +62,13 @@ def generate(
 
     formatted_prompt = format_prompt(prompt, history)
 
-    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+    
+    encoded = tokenizer(formatted_prompt, return_tensors="pt", add_special_tokens=False)
+    model_input = encoded
+    model.to(device)
+    generated_ids = model.generate(**model_input, max_new_tokens=1048, do_sample=True)
+    stream = tokenizer.batch_decode(generated_ids)
+    
     output = ""
 
     for response in stream: