Update app.py

app.py

@@ -1,7 +1,8 @@
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
 from peft import PeftModel
 import gradio as gr
+import threading
 
 # --- Load Model & Tokenizer ---
 
@@ -30,24 +31,48 @@ model = PeftModel.from_pretrained(model, adapter_model_name)
 model = model.merge_and_unload()
 print("Model ready!")
 
-# --- Gradio Function ---
+# --- Gradio Streaming Function ---
 
-def generate_text(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=100)
-    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return text
-
-# --- Gradio Interface ---
-def chat_fn(message, history):
+def chat_fn(message, history, max_tokens):
     prompt = message
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=100)
-    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return text
+
+    streamer = TextIteratorStreamer(
+        tokenizer,
+        skip_special_tokens=True
+    )
+
+    generation_kwargs = dict(
+        **inputs,
+        streamer=streamer,
+        max_new_tokens=int(max_tokens),
+        temperature=0.7,
+    )
+
+    thread = threading.Thread(
+        target=model.generate,
+        kwargs=generation_kwargs
+    )
+    thread.start()
+
+    partial_text = ""
+    for new_text in streamer:
+        partial_text += new_text
+        yield partial_text
+
+# --- Create Dropdown Component for max tokens ---
+
+dropdown = gr.Dropdown(
+    choices=["100", "200", "300"],
+    value="100",
+    label="Max New Tokens"
+)
+
+# --- Launch Gradio Chat Interface ---
 
 gr.ChatInterface(
     fn=chat_fn,
+    additional_inputs=[dropdown],
     title="LLM Finetuned Comment Generator",
     description="Chat with the model.",
 ).launch(share=True)