Add text streamer

app.py

@@ -4,11 +4,12 @@ import pandas as pd
 import numpy as np
 import torch
 import subprocess
+from threading import Thread
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     AutoProcessor,
-    TextStreamer
+    TextIteratorStreamer
 )
 
 # ─── MODEL SETUP ────────────────────────────────────────────────────────────────
@@ -99,9 +100,10 @@ def preview_csv(csv_file):
     
     # Create plot with first column as default
     first_column = column_choices[0]
+    df_with_index["_internal_idx"] = np.arange(len(df[first_column].values))
     plot = gr.LinePlot(
         df_with_index,
-        # x="index",
+        x="_internal_idx",
         y=first_column,
         title=f"Time Series: {first_column}"
     )
@@ -127,10 +129,11 @@ def update_plot(csv_file, selected_column):
         return gr.LinePlot(value=pd.DataFrame())
     
     df_with_index = df.copy()
+    df_with_index["_internal_idx"] = np.arange(len(df[selected_column].values))
     
     plot = gr.LinePlot(
         df_with_index,
-        # x="index",
+        x="_internal_idx",
         y=selected_column,
         title=f"Time Series: {selected_column}"
     )
@@ -190,30 +193,36 @@ def infer_chatts_stream(prompt: str, csv_file):
         inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
         # Generate with streaming
-        # streamer = TextStreamer(tokenizer)
-        generated_text = ""
-        with torch.no_grad():
-            outputs = model.generate(
-                **inputs, 
-                max_new_tokens=512,
-                do_sample=True,
-                temperature=0.7,
-                pad_token_id=tokenizer.eos_token_id
-            )
-            
-            # Decode the generated text
-            full_generated = tokenizer.decode(
-                outputs[0][inputs["input_ids"].shape[-1]:],
-                skip_special_tokens=True
-            )
-            
-            # Simulate streaming by yielding character by character
-            for i, char in enumerate(full_generated):
-                generated_text += char
-                if i % 5 == 0:  # Update every 5 characters for smoother streaming
-                    yield generated_text
+        streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
+        inputs.update({
+            "max_new_tokens": 512,
+            "streamer": streamer,
+            "temperature": 0.3
+        })
+        thread = threading.Thread(
+            target=model.generate,
+            kwargs=inputs
+        )
+        thread.start()
+
+        model_output = ""
+        for new_text in streamer:
+            model_output += new_text
+            yield model_output
+        
+        # # Decode the generated text
+        # full_generated = tokenizer.decode(
+        #     outputs[0][inputs["input_ids"].shape[-1]:],
+        #     skip_special_tokens=True
+        # )
+        
+        # # Simulate streaming by yielding character by character
+        # for i, char in enumerate(full_generated):
+        #     generated_text += char
+        #     if i % 5 == 0:  # Update every 5 characters for smoother streaming
+        #         yield generated_text
                     
-        yield generated_text
+        # yield generated_text
         
     except Exception as e:
         yield f"Error during inference: {str(e)}"