Clean 700M

app.py

@@ -1,31 +1,32 @@
+import gradio as gr
 import torch
 from transformers import pipeline
 
-# Set up the text-generation pipeline
 model_name = "amusktweewt/tiny-model-700M-chat"
+
 chatbot = pipeline(
     "text-generation",
     model=model_name,
     device=0 if torch.cuda.is_available() else -1
 )
 
-# Ensure that bos_token and eos_token are explicitly set as strings
-chatbot.tokenizer.bos_token = "<sos>"
-chatbot.tokenizer.eos_token = "<|endoftext|>"
-
-system_prompt = "You are a highly intelligent and helpful AI assistant named Tiny Chat, developed by amusktweewt. Always refer to yourself like that. Your responses should be clear, concise, and accurate. Always prioritize user needs, provide well-structured answers, and maintain a friendly yet professional tone. Adapt to the user's preferences and communication style. When needed, ask clarifying questions to ensure the best response. Be honest about limitations and avoid making assumptions. Keep interactions engaging, informative, and efficient."})
+system_prompt_default = (
+    "You are a highly intelligent and helpful AI assistant named Tiny Chat, "
+    "developed by amusktweewt. Always refer to yourself like that. "
+    "Your responses should be clear, concise, and accurate. "
+    "Always prioritize user needs, provide well-structured answers, "
+    "and maintain a friendly yet professional tone. "
+    "Adapt to the user's preferences and communication style."
+)
 
-def chat_fxn_caller(message, history, system_prompt="", temperature=0.6, top_p=0.95, max_tokens=32768):
+def chat_fxn_caller(message, history, system_prompt=system_prompt_default, temperature=0.6, top_p=0.95, max_tokens=512):
     messages = []
-
     if system_prompt.strip():
         messages.append({"role": "system", "content": system_prompt})
-
     for human_msg, assistant_msg in history:
         messages.append({"role": "user", "content": human_msg})
         if assistant_msg:
             messages.append({"role": "assistant", "content": assistant_msg})
-
     messages.append({"role": "user", "content": message})
 
     prompt = chatbot.tokenizer.apply_chat_template(messages, tokenize=False)
@@ -33,32 +34,23 @@ def chat_fxn_caller(message, history, system_prompt="", temperature=0.6, top_p=0
     response = chatbot(
         prompt,
         do_sample=True,
-        max_new_tokens=512,
+        max_new_tokens=max_tokens,
         top_k=50,
-        temperature=0.6,
+        temperature=temperature,
         num_return_sequences=1,
         repetition_penalty=1.1,
-        pad_token_id=chatbot.tokenizer.eos_token_id,
-        min_new_tokens=20
+        pad_token_id=chatbot.tokenizer.eos_token_id
     )
 
-    full_text = response[0]["generated_text"]
-    response = full_text[len(demo = gr.ChatInterface(
-    chat_fxn_caller,
-    type="messages",
-    additional_inputs=[
-        gr.Textbox(prompt, label="System Prompt"),
-    ],
-)
-
-demo.launch(share=True)prompt):].strip()
-    return response
+    # Extract only the new text after the prompt
+    generated = response[0]["generated_text"][len(prompt):].strip()
+    return generated
 
 demo = gr.ChatInterface(
     chat_fxn_caller,
     type="messages",
     additional_inputs=[
-        gr.Textbox(system_prompt, label="System Prompt"),
+        gr.Textbox(system_prompt_default, label="System Prompt"),
     ],
 )