Update app.py

app.py

@@ -1,6 +1,7 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import time
+import gradio as gr
 
 def generate_prompt(instruction, input=""):
     instruction = instruction.strip().replace('\r\n','\n').replace('\n\n','\n')
@@ -25,10 +26,10 @@ model_path = "models/rwkv-6-world-1b6/" # Path to your local model directory
 model = AutoModelForCausalLM.from_pretrained(
     model_path, 
     trust_remote_code=True, 
-    use_flash_attention_2=False  # Explicitly disable Flash Attention
+    use_flash_attention_2=False
 ).to(torch.float32)
 
-
+# Create a custom tokenizer (make sure to download vocab.json)
 tokenizer = AutoTokenizer.from_pretrained(
     model_path, 
     bos_token="</s>",
@@ -40,23 +41,41 @@ tokenizer = AutoTokenizer.from_pretrained(
     clean_up_tokenization_spaces=False  # Or set to True if you prefer
 )
 
-print(tokenizer.special_tokens_map)
-
-text = "Hi"
-
-prompt = generate_prompt(text)
-
-input_ids = tokenizer(prompt, return_tensors="pt").input_ids
-
-# Generate text word by word with stop sequence
-generated_text = ""
-for i in range(333):  # Generate up to 333 tokens
-    output = model.generate(input_ids, max_new_tokens=1, do_sample=True, temperature=1.0, top_p=0.3, top_k=0)
-    new_word = tokenizer.decode(output[0][-1:], skip_special_tokens=True)
-
-    print(new_word, end="", flush=True)  # Print word-by-word
-    generated_text += new_word
+# Function to handle text generation with word-by-word output and stop sequence
+def generate_text(input_text):
+    prompt = generate_prompt(input_text)
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+
+    generated_text = ""
+    stop_sequence_found = False
+    for i in range(333):
+        output = model.generate(input_ids, max_new_tokens=1, do_sample=True, temperature=1.0, top_p=0.3, top_k=0)
+        new_word = tokenizer.decode(output[0][-1:], skip_special_tokens=True)
+
+        print(new_word, end="", flush=True)
+        generated_text += new_word
+
+        if new_word == '\n' or new_word == '.':
+            stop_sequence_found = True
+            break
+
+        input_ids = output
+
+    if stop_sequence_found:
+        print("\n(Stop sequence found)")
+    print()
+    return generated_text
+
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=generate_text,
+    inputs="text",
+    outputs="text",
+    title="RWKV Chatbot",
+    description="Enter your prompt below:",
+)
 
-    input_ids = output  # Update input_ids for next iteration
+# For local testing:
+# iface.launch()
 
-print()  # Add a newline at the end
+# Hugging Face Spaces will automatically launch the interface.