Update app.py

app.py

@@ -1,44 +1,53 @@
 import gradio as gr
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoModelForCausalLM
 import torch
 
-# Load the model and tokenizer
+# Load the model
 model_name = "wop/kosmox-gguf"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModel.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# Define the chat template function
+def format_chat(messages, add_generation_prompt):
+    formatted = "<BOS>"
+    for message in messages:
+        if message['from'] == 'human':
+            formatted += ' ' + message['value'] + ' '
+        elif message['from'] == 'gpt':
+            formatted += ' ' + message['value'] + ' '
+        else:
+            formatted += '<|' + message['from'] + '|> ' + message['value'] + ' '
+    if add_generation_prompt:
+        formatted += ' '
+    return formatted
 
 # Function to generate responses
 def respond(message, history, system_message, max_tokens, temperature, top_p):
     # Prepare the chat history
-    messages = [{"role": "system", "content": system_message}]
-    
+    messages = [{"from": "system", "value": system_message}]
     for user_msg, bot_msg in history:
         if user_msg:
-            messages.append({"role": "user", "content": user_msg})
+            messages.append({"from": "human", "value": user_msg})
         if bot_msg:
-            messages.append({"role": "assistant", "content": bot_msg})
-    
-    messages.append({"role": "user", "content": message})
+            messages.append({"from": "gpt", "value": bot_msg})
+    messages.append({"from": "human", "value": message})
 
-    # Create the chat input for the model
-    chat_input = tokenizer.chat_template.format(
-        bos_token=tokenizer.bos_token,
-        messages=[{"from": "human", "value": m['content']} if m['role'] == 'user' else {"from": "gpt", "value": m['content']} for m in messages]
-    )
-    
-    inputs = tokenizer(chat_input, return_tensors="pt")
+    # Format the chat input for the model
+    chat_input = format_chat(messages, add_generation_prompt=False)
+
+    # Tokenize input (assuming model can handle raw text inputs internally)
+    inputs = torch.tensor([ord(c) for c in chat_input]).unsqueeze(0)  # Dummy tokenization
 
     # Generate response
     with torch.no_grad():
         outputs = model.generate(
-            input_ids=inputs['input_ids'],
+            input_ids=inputs,
             max_length=max_tokens,
             temperature=temperature,
             top_p=top_p,
             do_sample=True
         )
     
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    response = ''.join([chr(t) for t in outputs[0].tolist() if t < 256])  # Dummy decoding
     yield response.strip()
 
 # Define the Gradio interface