Update app.py

app.py

@@ -4,8 +4,10 @@ import torch
 import torch.nn as nn
 import torchvision.transforms as transforms
 import torchvision.models as models
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import os
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
+import torch
+import gc
 
 # Set device
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -66,27 +68,70 @@ preprocess = transforms.Compose([
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # ImageNet normalization
 ])
 
+
 # Load Meta's LLaMA model for generating product descriptions
 def load_llama():
     model_name = "meta-llama/Llama-3.2-1B-Instruct"  
     token = os.getenv("HUGGINGFACE_TOKEN") 
     tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=token)
     model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=token).to(device)
+     # Initialize the text generation pipeline with the prepared model
+    text_generation = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer
+    )
     return tokenizer, model
-
 llama_tokenizer, llama_model = load_llama()
 
-# Generate product description using LLaMA
+# Generate product description using external data and structured format
 def generate_description(category, subclass):
-    prompt = f"Generate a detailed and engaging product description for a {category} of type {subclass}."
-    
-    inputs = llama_tokenizer.encode(prompt, return_tensors="pt").to(device)
-    outputs = llama_model.generate(inputs, max_length=100, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
-    description = llama_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Define file path and read content
+    file_path = 'data for product description.txt'
+    with open(file_path, 'r', encoding='utf-8') as file:
+        file_content = file.read()
+
+    prompt = f"""
+    [Data]
+    {file_content}
+    Role: You are a product description content writer with 10 years of experience in the market. Generate a product description for a {subclass} in the {category} category based on the [Data] provided.
+    Follow the [Instructions] strictly:
+    [Instructions]
+    - Create a detailed product description for a {subclass} in the {category} category based on the [Data].
+    - Use the structured format below, making each section clear and concise.
+    - Highlight key product features, technical specifications, and the target audience.
+    """
+
+    generated_texts = llama_model.generate(
+        inputs=llama_tokenizer(prompt, return_tensors="pt").input_ids.to(device),
+        max_length=7000,
+        max_new_tokens=2000,
+        do_sample=True,
+        temperature=0.7,
+        top_k=50,
+        top_p=0.95,
+    )
+
+    description = llama_tokenizer.decode(generated_texts[0], skip_special_tokens=True)
     
+    # Clean up resources
+    torch.cuda.empty_cache()
+    gc.collect()
+
     return description
 
 
+# # Generate product description using LLaMA
+# def generate_description(category, subclass):
+#     prompt = f"Generate a detailed and engaging product description for a {category} of type {subclass}."
+    
+#     inputs = llama_tokenizer.encode(prompt, return_tensors="pt").to(device)
+#     outputs = llama_model.generate(inputs, max_length=100, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
+#     description = llama_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+#     return description
+
+
 def classify_image(image):
     # Open the image using PIL
     image = Image.fromarray(image)