Update src/text_processor.py

src/text_processor.py

@@ -1,20 +1,11 @@
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+# --- src/text_processor.py ---
+from huggingface_hub import InferenceClient
 from pydantic import BaseModel
-import spaces
 
-torch.random.manual_seed(0)
-model = AutoModelForCausalLM.from_pretrained(
+# Hugging Face Hub client setup
+client = InferenceClient(
     "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,
+    token = api_key=os.environ['HUGGINGFACE_API_KEY']    
 )
 
 # Pydantic class for output validation
@@ -24,13 +15,11 @@ 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:
@@ -38,21 +27,20 @@ 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}} 
     """
 
-    generation_args = {
-        "max_new_tokens": 100,  # Adjust as needed
-        "return_full_text": False,
-        "temperature": 0.0,
-        "do_sample": False,
-    }
+    # 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
 
-    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)