edit

app.py

@@ -3,64 +3,53 @@ import torch
 from unsloth import FastLanguageModel
 from transformers import AutoTokenizer
 
-# Load your fine-tuned Phi-3 model
-def load_phi3():
-    model_name = "jcrissa/phi3-new-t2i"  # Your trained model
-    max_seq_length = 4096  # Ensure correct max length
-    
-    # Load fine-tuned model
+# Load your fine-tuned Phi-3 model from Hugging Face
+MODEL_NAME = "jcrissa/phi3-new-t2i"
+
+def load_phi3_model():
     model, tokenizer = FastLanguageModel.from_pretrained(
-        model_name,
-        max_seq_length=max_seq_length,
-        dtype=None,  # Uses default torch dtype (float16 or bfloat16 if available)
-        load_in_4bit=True  # Uses 4-bit quantization for efficiency
+        MODEL_NAME,
+        max_seq_length=4096,  # Ensure it matches your fine-tuning
+        dtype=None  # Use `torch.float16` if running on GPU
     )
     
-    tokenizer.pad_token = tokenizer.eos_token  # Ensure padding is set
+    tokenizer.pad_token = tokenizer.eos_token
     tokenizer.padding_side = "left"
-
+    
     return model, tokenizer
 
-# Initialize model and tokenizer
-phi3_model, phi3_tokenizer = load_phi3()
+phi3_model, phi3_tokenizer = load_phi3_model()
 
-# Function to generate prompts
+# Function to generate text using Phi-3
 def generate(plain_text):
-    input_ids = phi3_tokenizer(plain_text.strip(), return_tensors="pt").input_ids.cuda()  # Move to GPU if available
+    input_ids = phi3_tokenizer(plain_text.strip(), return_tensors="pt").input_ids
     eos_id = phi3_tokenizer.eos_token_id
     
     outputs = phi3_model.generate(
         input_ids, 
-        do_sample=True, 
+        do_sample=True,  
         max_new_tokens=75, 
-        num_beams=5, 
+        num_beams=8, 
         num_return_sequences=1, 
         eos_token_id=eos_id, 
-        pad_token_id=eos_id, 
-        length_penalty=1.0
+        pad_token_id=eos_id,
+        length_penalty=-1.0
     )
     
-    output_texts = phi3_tokenizer.batch_decode(outputs, skip_special_tokens=True)
-    return output_texts[0]
+    output_text = phi3_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return output_text.strip()
 
-# Gradio UI
-txt = grad.Textbox(lines=1, label="Initial Text", placeholder="Input Prompt")
-out = grad.Textbox(lines=1, label="Optimized Prompt")
-examples = [
-    "A rabbit is wearing a space suit", 
-    "Several railroad tracks with one train passing by", 
-    "The roof is wet from the rain", 
-    "Cats dancing in a space club"
-]
+# Setup Gradio Interface
+txt = grad.Textbox(lines=1, label="Input Text", placeholder="Enter your prompt")
+out = grad.Textbox(lines=1, label="Generated Text")
 
 grad.Interface(
     fn=generate,
     inputs=txt,
     outputs=out,
-    title="Phi-3 Prompt Generator",
-    description="Fine-tuned Phi-3 model (`jcrissa/phi3-new-t2i`) for text-to-image prompt generation.",
-    examples=examples,
-    allow_flagging='never',
+    title="Fine-Tuned Phi-3 Model",
+    description="This demo uses a fine-tuned Phi-3 model to optimize text prompts.",
+    allow_flagging="never",
     cache_examples=False,
     theme="default"
 ).launch(enable_queue=True, debug=True)