Update main.py

main.py

@@ -11,6 +11,9 @@ import io
 from transformers import pipeline, CLIPProcessor, CLIPModel
 from datetime import datetime
 import uvicorn
+# Hugging Face GPT or LLM model for content-based name generation
+from langchain.chat_models import ChatOpenAI
+import os
 
 app = FastAPI()
 
@@ -22,8 +25,19 @@ MAX_CONTENT_LENGTH = 16 * 1024 * 1024  # 16 MB
 if not os.path.exists(UPLOAD_FOLDER):
     os.makedirs(UPLOAD_FOLDER)
 
-# Hugging Face GPT or LLM model for content-based name generation
-llm_pipeline = pipeline("text-generation", model="gpt2")
+# Load your OpenAI API key from environment variables
+openai_api_key = os.getenv("OPENAI_API_KEY")
+
+# Ensure the API key is correctly loaded
+if openai_api_key is None:
+    raise ValueError("API key not found. Please set your OPENAI_API_KEY environment variable.")
+
+# Initialize the LLM (Language Model) with GPT-4o-mini or other available model
+llm = ChatOpenAI(
+    model_name="gpt-4o-mini",  # Specify the correct model name (e.g., "gpt-4" or "gpt-4o-mini")
+    temperature=0,  # Set temperature to 0 for deterministic responses (no randomness)
+    openai_api_key=openai_api_key  # Pass the OpenAI API key
+)
 
 # Load the CLIP model for image feature extraction
 clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
@@ -31,9 +45,12 @@ clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
 # Function to generate a more appropriate name based on content
 def generate_name_based_on_content(text):
-    prompt = f"Generate a meaningful file name for the following content: {text[:200]}"
-    result = llm_pipeline(prompt, max_length=10, num_return_sequences=1)
-    return result[0]['generated_text'].strip()
+    prompt = f"Generate a meaningful file name for the following content: {text[:200]}"  # Truncate text to first 200 characters
+    response = llm(prompt)  # Get the model's response
+
+    # Extract the generated file name and clean it
+    file_name = response.strip()  # Strip any unnecessary whitespace or characters
+    return file_name
 
     
 # Allowed file extensions check