Fixed model not working

- Not an ideal solution, streaming doesn't work

app.py

@@ -1,5 +1,9 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
+from unsloth import FastLanguageModel
+from unsloth.chat_templates import get_chat_template
+from transformers import TextStreamer
+
 
 """
 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
@@ -7,8 +11,20 @@ For more information on `huggingface_hub` Inference API support, please check th
 
 model_name_or_path = "samlama111/lora_model"
 
-client = InferenceClient(model_name_or_path)
+# client = InferenceClient(model_name_or_path)
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name = model_name_or_path,
+    max_seq_length = 8192,
+    load_in_4bit = True,
+    # token = "hf_...", # No need since our model is public
+)   
 
+tokenizer = get_chat_template(
+    tokenizer,
+    chat_template = "llama-3.1",
+    mapping = {"role" : "from", "content" : "value", "user" : "human", "assistant" : "gpt"}, # ShareGPT style
+)
+FastLanguageModel.for_inference(model) # Enable native 2x faster inference
 
 def respond(
     message,
@@ -29,17 +45,19 @@ def respond(
     messages.append({"role": "user", "content": message})
 
     response = ""
+    
+    inputs = tokenizer.apply_chat_template(messages, tokenize = True, add_generation_prompt = True, return_tensors = "pt")
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+    text_streamer = TextStreamer(tokenizer)
+    # TODO: Doesn't stream ATM
+    for message in model.generate(input_ids = inputs, streamer = text_streamer, max_new_tokens = 1024, use_cache = True):
+        # Decode the tensor to a string
+        decoded_message = tokenizer.decode(message, skip_special_tokens=True)
+        
+        # Manually getting the response
+        response = decoded_message.split("assistant")[-1].strip()  # Extract only the assistant's response
+        print(response)
 
-        response += token
         yield response