Update app.py

app.py

@@ -1,34 +1,33 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline
 import torch
 
-# Load the DeepSeek-R1-Distill-Qwen-1.5B-uncensored model
-model_id = "thirdeyeai/DeepSeek-R1-Distill-Qwen-1.5B-uncensored"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    torch_dtype=torch.float16,  # Use float16 for efficiency
-    low_cpu_mem_usage=True,
-    device_map="auto"  # Automatically use available devices
+# Initialize the text generation pipeline with the model
+generator = pipeline(
+    "text-generation", 
+    model="thirdeyeai/DeepSeek-R1-Distill-Qwen-1.5B-uncensored",
+    torch_dtype=torch.float16,
+    device_map="auto"
 )
 
 def generate_text(prompt, max_length=100, temperature=0.7, top_p=0.9):
-    """Generate text based on prompt"""
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    """Generate text based on prompt using the pipeline"""
+    # Calculate max_new_tokens from max_length
+    # This is approximate as token count doesn't directly map to character count
+    max_new_tokens = max_length // 4  # rough estimate of 4 chars per token
     
-    # Generate
-    with torch.no_grad():
-        generation_output = model.generate(
-            input_ids=inputs.input_ids,
-            attention_mask=inputs.attention_mask,
-            max_length=len(inputs.input_ids[0]) + max_length,
-            temperature=temperature,
-            top_p=top_p,
-            do_sample=True,
-        )
+    # Generate text
+    response = generator(
+        prompt,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        do_sample=True,
+        return_full_text=True
+    )
     
-    # Decode and return only the generated part
-    generated_text = tokenizer.decode(generation_output[0], skip_special_tokens=True)
+    # Extract the generated text from the response
+    generated_text = response[0]['generated_text']
     return generated_text
 
 # Create Gradio interface
@@ -36,7 +35,7 @@ demo = gr.Interface(
     fn=generate_text,
     inputs=[
         gr.Textbox(lines=5, placeholder="Enter your prompt here...", label="Prompt"),
-        gr.Slider(minimum=10, maximum=500, value=100, step=10, label="Max Length"),
+        gr.Slider(minimum=10, maximum=500, value=100, step=10, label="Max Length (approx. characters)"),
         gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p")
     ],