Create text_processor.py

src/text_processor.py

@@ -0,0 +1,62 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from pydantic import BaseModel
+import spaces
+
+torch.random.manual_seed(0)
+model = AutoModelForCausalLM.from_pretrained(
+    "microsoft/Phi-3-mini-4k-instruct",
+    device_map="cuda",
+    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
+class VideoAnalysis(BaseModel):
+    indoor: int
+    hands_free: int
+    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:
+    - Is the scene indoors? 
+    - 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}} 
+    """
+
+    generation_args = {
+        "max_new_tokens": 100,  # Adjust as needed
+        "return_full_text": False,
+        "temperature": 0.0,
+        "do_sample": False,
+    }
+
+    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."}