Update src/text_processor.py

src/text_processor.py

@@ -1,12 +1,23 @@
-# --- src/text_processor.py ---
-from huggingface_hub import InferenceClient
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from pydantic import BaseModel
-import os
+import spaces 
 
-# Hugging Face Hub client setup
-client = InferenceClient(
+device = 'cuda'
+
+# Load your LLM model and tokenizer 
+torch.random.manual_seed(0)
+model = AutoModelForCausalLM.from_pretrained(
     "microsoft/Phi-3-mini-4k-instruct",
-    token = os.environ['HUGGINGFACE_API_KEY']    
+    device_map=device,
+    torch_dtype="auto",
+    trust_remote_code=True,
+)
+tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct")
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
 )
 
 # Pydantic class for output validation
@@ -16,11 +27,12 @@ class VideoAnalysis(BaseModel):
     screen_interaction: int
     standing: int
 
+@spaces.GPU(duration=100)
 def process_description(description):
     # Construct a prompt for your LLM based on the video description
     prompt = f"""
     You are a helpful AI assistant. Analyze the following video description and answer the questions with 0 for True and 1 for False: 
-
+    
     Video Description: {description}
 
     Questions:
@@ -28,18 +40,25 @@ def process_description(description):
     - Are the subject's hands free?
     - Is there screen interaction by the subject?
     - Is the subject standing? 
-
+    
     Provide your answers in JSON format like this: 
     {{"indoor": 0, "hands_free": 1, "screen_interaction": 0, "standing": 1}} 
     """
 
-    # Using the Hugging Face Hub InferenceClient for text generation
-    response = client.chat_completion(
-        messages=[{"role": "user", "content": prompt}],
-        max_tokens=100,  # Adjust as needed
-    )
-
-    # Extract the generated JSON text from the response
-    json_text = response.choices[0].message.content
+    generation_args = {
+        "max_new_tokens": 100,  # Adjust as needed
+        "return_full_text": False,
+        "temperature": 0.0,
+        "do_sample": False,
+    }
 
-    return json_text
+    output = pipe(prompt, **generation_args)
+    json_text = output[0]['generated_text']
+    
+    try:
+        # Attempt to parse and validate the JSON response
+        analysis_result = VideoAnalysis.model_validate_json(json_text)
+        return analysis_result.model_dump_json()  # Return as valid JSON
+    except Exception as e:
+        print(f"Error processing LLM output: {e}")
+        return {"error": "Could not process the video description."}