Update app.py

app.py

@@ -1,13 +1,34 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import onnxruntime as ort
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
 
-# Load the Phi-3.5-mini-instruct model and tokenizer
+# Set the random seed for reproducibility
+torch.random.manual_seed(0)
+
+# Load the model and tokenizer
 model_name = "microsoft/Phi-3.5-mini-instruct"
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, 
+    device_map="auto", 
+    torch_dtype="auto", 
+    trust_remote_code=True
+)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Load the ONNX model
-session = ort.InferenceSession(f"{model_name}/model.onnx")
+# Set up the pipeline
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+)
+
+# Define the generation arguments
+generation_args = {
+    "max_new_tokens": 150,
+    "return_full_text": False,
+    "temperature": 0.7,
+    "do_sample": False,
+}
 
 # Simple HTML template for the website
 simple_website_template = """
@@ -41,22 +62,21 @@ simple_website_template = """
 
 # Function to generate personalized content using Phi-3.5-mini-instruct
 def personalize_website_llm(persona_text):
-    # Create a prompt for the model
-    prompt = f"Generate personalized website content for the following persona: {persona_text}. Provide a title and main content."
+    # Construct the conversation history
+    messages = [
+        {"role": "system", "content": "You are a helpful AI assistant that personalizes content for websites."},
+        {"role": "user", "content": f"Persona: {persona_text}. Generate a personalized website content including a title and a paragraph."},
+    ]
     
-    # Tokenize the prompt
-    inputs = tokenizer(prompt, return_tensors="np")
+    # Generate content using the pipeline
+    output = pipe(messages, **generation_args)
+    generated_text = output[0]['generated_text'].strip()
+
+    # Simple heuristic to split title and content
+    lines = generated_text.split('\n')
+    title = lines[0]
+    content = "\n".join(lines[1:])
 
-    # Run the ONNX model
-    ort_inputs = {session.get_inputs()[0].name: inputs["input_ids"]}
-    ort_outs = session.run(None, ort_inputs)
-    
-    # Decode the output
-    generated_text = tokenizer.decode(ort_outs[0][0], skip_special_tokens=True)
-    
-    # Split the response into a title and content
-    title, content = generated_text.split('\n', 1)
-    
     # Set the title color and font size based on simple heuristics
     title_color = "#333"
     font_size = 16
@@ -73,8 +93,8 @@ def personalize_website_llm(persona_text):
     personalized_website = simple_website_template.format(
         title_color=title_color,
         font_size=font_size,
-        title=title.strip(),
-        content=content.strip()
+        title=title,
+        content=content
     )
     
     return personalized_website