Streaming enabled

app.py

@@ -14,8 +14,6 @@ st.markdown(
     """
 )
 
-
-
 hf_token = st.sidebar.text_input("Enter your Hugging Face API Token", type="password")
 
 # -------------------------------
@@ -49,14 +47,12 @@ def get_device():
 @st.cache_resource(show_spinner=True)
 def load_model(hf_token):
     device = get_device()
-    # Load the gated base model with your token
     base_model = AutoModelForCausalLM.from_pretrained(
         "meta-llama/Llama-3.2-1B-Instruct",
         trust_remote_code=True,
         torch_dtype=torch.float16,
         use_auth_token=hf_token
     )
-    # Load the tokenizer from your adapter repository and set pad token if needed
     tokenizer = AutoTokenizer.from_pretrained(
         "Akash190104/space_turtle_101",
         use_fast=False,
@@ -65,7 +61,6 @@ def load_model(hf_token):
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
-    # Load the LoRA adapter using your token
     model = PeftModel.from_pretrained(
         base_model,
         "Akash190104/space_turtle_101",
@@ -83,23 +78,27 @@ else:
             st.success("Model loaded successfully!")
         except Exception as e:
             st.error(f"Model loading failed: {e}")
-            st.error("Ensure your token has access to meta-llama/Llama-3.2-1B-Instruct.")
             st.stop()
 
 # -------------------------------
-# Generation helper function
+# Streaming Generation Function
 # -------------------------------
-def generate_sample(prompt_text):
+def stream_generate(prompt_text):
     inputs = tokenizer(prompt_text, return_tensors="pt", padding=True).to(device)
+    result_placeholder = st.empty()  # Placeholder for streaming text
+    generated_text = ""
     with torch.inference_mode():
         outputs = model.generate(
             **inputs,
             max_new_tokens=150,
             do_sample=True,
             temperature=0.7,
-            top_p=0.9
+            top_p=0.9,
+            streamer=True  # Enables token-by-token streaming
         )
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        for token in outputs:
+            generated_text += tokenizer.decode([token], skip_special_tokens=True)
+            result_placeholder.text(generated_text)  # Update in real-time
     return generated_text
 
 # -------------------------------
@@ -133,10 +132,11 @@ if mode == "Interactive":
             st.error("Please provide both a bias category and a country/region.")
         else:
             prompt = f"```{bias_input} in {country_input}```\n"
-            generated = generate_sample(prompt)
+            generated = stream_generate(prompt)
             st.markdown("**Generated Output:**")
             st.text_area("", value=generated, height=200)
             st.download_button("Download Output", generated, file_name="output.txt")
+
 elif mode == "Random Generation (10 samples)":
     st.subheader("Random Generation Mode")
     if st.button("Generate 10 Random Samples"):
@@ -145,7 +145,7 @@ elif mode == "Random Generation (10 samples)":
             bias = random.choice(biases)
             country = random.choice(countries)
             prompt = f"```{bias} in {country}```\n"
-            generated = generate_sample(prompt)
+            generated = stream_generate(prompt)
             results.append({"prompt": prompt, "generated": generated})
         for i, res in enumerate(results):
             st.markdown(f"**Sample {i+1}:**")