Update app.py

app.py

@@ -1,45 +1,46 @@
-# import gradio as gr
+from flask import Flask, request, jsonify
 import os
-import torch
-from transformers import AutoProcessor, MllamaForConditionalGeneration, TextIteratorStreamer
+import requests
 from PIL import Image
-import spaces
 import tempfile
-import requests
 from PyPDF2 import PdfReader
 from threading import Thread
-from flask import Flask, request, jsonify
 import io
 import fitz
+from groq import Groq
+from queue import Queue
 
-# Check if we're running in a Hugging Face Space and if SPACES_ZERO_GPU is enabled
-# IS_SPACES_ZERO = os.environ.get("SPACES_ZERO_GPU", "0") == "1"
-# IS_SPACE = os.environ.get("SPACE_ID", None) is not None
-
-# Determine the device (GPU if available, else CPU)
-device = "cuda" if torch.cuda.is_available() else "cpu"
-LOW_MEMORY = os.getenv("LOW_MEMORY", "0") == "1"
-
-print(f"Using device: {device}")
-print(f"Low memory mode: {LOW_MEMORY}")
-
+# Initialize Flask app
 app = Flask(__name__)
 
-# Get Hugging Face token from environment variables
+# Get API tokens from environment variables
 HF_TOKEN = os.environ.get('HF_TOKEN')
+GROQ_API_KEY = os.environ.get('GROQ_API_KEY')
 
-# Load the model and processor
-model_name = "meta-llama/Llama-3.2-11B-Vision-Instruct"
-model = MllamaForConditionalGeneration.from_pretrained(
-    model_name,
-    use_auth_token=HF_TOKEN,
-    torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
-    device_map="auto" if device == "cuda" else None,  # Use device mapping if CUDA is available
-)
+# Initialize Groq client
+client = Groq(api_key=GROQ_API_KEY)
 
-# Move the model to the appropriate device (GPU if available)
-# model.to(device)
-processor = AutoProcessor.from_pretrained(model_name, use_auth_token=HF_TOKEN)
+# Configuration for low memory mode (maintaining original functionality)
+LOW_MEMORY = os.getenv("LOW_MEMORY", "0") == "1"
+print(f"Low memory mode: {LOW_MEMORY}")
+
+class TextStreamer:
+    def __init__(self):
+        self.queue = Queue()
+        self.buffer = ""
+        
+    def put(self, text):
+        self.queue.put(text)
+        
+    def __iter__(self):
+        while True:
+            if not self.queue.empty():
+                text = self.queue.get()
+                if text is None:  # End signal
+                    break
+                yield text
+            else:
+                continue
 
 def extract_image_from_pdf(pdf_url, dpi=75):
     """
@@ -76,62 +77,62 @@ def extract_image_from_pdf(pdf_url, dpi=75):
         print(f"Error extracting first page: {e}")
         return None
 
-
-
 def predict_image(image_url, text, file_pref):
     try:
-        # Download the image from the URL
-        # response = requests.get(image_url)
-        # response.raise_for_status()  # Raise an error for invalid responses
-        # image = Image.open(io.BytesIO(response.content)).convert("RGB")
         if file_pref == 'img':
             response = requests.get(image_url)
-            response.raise_for_status()  # Raise an error for invalid responses
+            response.raise_for_status()
             image = Image.open(io.BytesIO(response.content)).convert("RGB")
         else: 
             image = extract_image_from_pdf(image_url)
             
-        messages = [
-            {"role": "user", "content": [
-                {"type": "image"},  # Specify that an image is provided
-                {"type": "text", "text": text}  # Add the user-provided text input
-            ]}
-        ]
+        streamer = TextStreamer()
         
-        # Create the input text using the processor's chat template
-        input_text = processor.apply_chat_template(messages, add_generation_prompt=True)
-        
-        # Process the inputs and move to the appropriate device
-        inputs = processor(image, input_text, return_tensors="pt").to(device)
-
-        # outputs = model.generate(**inputs, max_new_tokens=100)
-    
-        # # Decode the output to return the final response
-        # response = processor.decode(outputs[0], skip_special_tokens=True)
-               
-        # return response
-
-        streamer = TextIteratorStreamer(processor, skip_special_tokens=True, skip_prompt=True)
-
-        generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=4096)
-        generated_text = ""
-        
-        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        def generate_response():
+            try:
+                completion = client.chat.completions.create(
+                    model="mixtral-8x7b-32768",
+                    messages=[
+                        {
+                            "role": "user",
+                            "content": [
+                                {
+                                    "type": "image_url",
+                                    "image_url": {"url": image_url}
+                                },
+                                {
+                                    "type": "text",
+                                    "text": text
+                                }
+                            ]
+                        }
+                    ],
+                    temperature=0.7,
+                    max_tokens=4096,
+                    top_p=1,
+                    stream=True
+                )
+                
+                for chunk in completion:
+                    if chunk.choices[0].delta.content:
+                        streamer.put(chunk.choices[0].delta.content)
+                streamer.put(None)  # Signal the end
+            except Exception as e:
+                print(f"Error in generate_response: {e}")
+                streamer.put(None)
+
+        thread = Thread(target=generate_response)
         thread.start()
-        buffer = ""
         
+        buffer = ""
         for new_text in streamer:
             buffer += new_text
-            # generated_text_without_prompt = buffer
-            # # time.sleep(0.01)
-            # yield buffer
-        
+            
         return buffer
         
     except Exception as e:
         raise ValueError(f"Error during prediction: {str(e)}")
 
-
 def extract_text_from_pdf(pdf_url):
     try:
         response = requests.get(pdf_url)
@@ -149,48 +150,44 @@ def extract_text_from_pdf(pdf_url):
         return text
     except Exception as e:
         raise ValueError(f"Error extracting text from PDF: {str(e)}")
-# raise HTTPException(status_code=400, detail=f"Error extracting text from PDF: {str(e)}")
 
-@spaces.GPU
 def predict_text(text):
-    # pdf_text = extract_text_from_pdf('https://arinsight.co/2024_FA_AEC_1200_GR1_GR2.pdf')
-    
-    text_combined = text # + "\n\nExtracted Text from PDF:\n" + pdf_text
-    
-    # Prepare the input messages
-    messages = [{"role": "user", "content": [{"type": "text", "text": text_combined}]}]
+    streamer = TextStreamer()
     
-    # Create the input text using the processor's chat template
-    input_text = processor.apply_chat_template(messages, add_generation_prompt=True)
-    
-    # Process the inputs and move to the appropriate device
-    # inputs = processor(image, input_text, return_tensors="pt").to(device)
-    inputs = processor(text=input_text, return_tensors="pt").to("cuda")
-    # Generate a response from the model
-    # outputs = model.generate(**inputs, max_new_tokens=1024)
-    
-    # # Decode the output to return the final response
-    # response = processor.decode(outputs[0], skip_special_tokens=True, skip_prompt=True)
-
-
-    streamer = TextIteratorStreamer(processor, skip_special_tokens=True, skip_prompt=True)
-
-    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=2048)
-    generated_text = ""
-    
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    def generate_response():
+        try:
+            completion = client.chat.completions.create(
+                model="mixtral-8x7b-32768",
+                messages=[
+                    {
+                        "role": "user",
+                        "content": text
+                    }
+                ],
+                temperature=0.7,
+                max_tokens=2048,
+                top_p=1,
+                stream=True
+            )
+            
+            for chunk in completion:
+                if chunk.choices[0].delta.content:
+                    streamer.put(chunk.choices[0].delta.content)
+            streamer.put(None)  # Signal the end
+        except Exception as e:
+            print(f"Error in generate_response: {e}")
+            streamer.put(None)
+
+    thread = Thread(target=generate_response)
     thread.start()
-    buffer = ""
     
+    buffer = ""
     for new_text in streamer:
         buffer += new_text
-        # generated_text_without_prompt = buffer
-        # # time.sleep(0.01)
-        # yield buffer
-    
+        
     return buffer
 
-
+# [Rest of the prompts remain exactly the same as in original]
 PROMPT = (
     "Extract the following information as per this format:\n"
     "'Course Code:'\n"
@@ -213,47 +210,6 @@ PROMPT_SKILLS = (
     "'Tertiary Skills'."
 )
 
-
-# PROMPT_IMAGE = (
-#     "You are a highly intelligent assistant designed to analyze images and extract structured information from them. "
-#     "Your task is to analyze the given image of a student's academic record and generate a response in the exact JSON format provided below. "
-#     "If any specific information is missing or unavailable in the image, replace the corresponding field with null. "
-#     "Ensure the format is consistent, strictly adhering to the structure shown below.\n\n"
-#     "Required JSON Format:\n\n"
-#     "{\n"
-#     '  "student": {\n'
-#     '    "name": "string",\n'
-#     '    "id": "string",\n'
-#     '    "dob": "string",\n'
-#     '    "original_start_date": "string",\n'
-#     '    "cumulative_gpa": "string",\n'
-#     '    "program": "string",\n'
-#     '    "status": "string"\n'
-#     '  },\n'
-#     '  "courses": [\n'
-#     '    {\n'
-#     '      "transfer_institution": "string",\n'
-#     '      "course_code": "string",\n'
-#     '      "course_name": "string",\n'
-#     '      "credits_attempted": number,\n'
-#     '      "credits_earned": number,\n'
-#     '      "grade": "string",\n'
-#     '      "quality_points": number,\n'
-#     '      "semester_code": "string",\n'
-#     '      "semester_dates": "string"\n'
-#     '    }\n'
-#     "    // Additional courses can be added here\n"
-#     "  ]\n"
-#     "}\n\n"
-#     "Instructions:\n\n"
-#     "1. Extract the student information and course details as displayed in the image.\n"
-#     "2. Use null for any missing or unavailable information.\n"
-#     "3. Format the extracted data exactly as shown above. Do not deviate from this structure.\n"
-#     "4. Use accurate field names and ensure proper nesting of data (e.g., 'student' and 'courses' sections).\n"
-#     "5. The values for numeric fields like credits_attempted, credits_earned, and quality_points should be numbers (not strings).\n"
-# )
-
-
 PROMPT_IMAGE_STUDENT = (
     "You are a highly intelligent assistant designed to analyze images and extract structured information from them. "
     "Your task is to analyze the given image of a student's academic record and generate a response in the exact JSON format provided below. "
@@ -309,8 +265,6 @@ PROMPT_IMAGE_COURSES = (
     "5. Return only the 'courses' section as JSON.\n"
 )
 
-
-
 @app.route("/", methods=["GET"])
 def home():
     return jsonify({"message": "Welcome to the PDF Extraction API. Use the /extract endpoint to extract information."})
@@ -351,9 +305,6 @@ def extract_info():
             response_student = predict_image(img_url, prompt_student, file_pref)
             response_courses = predict_image(img_url, prompt_courses, file_pref)
             response_image = response_student + response_courses 
-
-            # response_image = {"student": response_student.get("student", {}), "courses": response_courses.get("courses", [])}
-            
         else:
             response_image = ''
             
@@ -363,21 +314,4 @@ def extract_info():
         return jsonify({"error": str(e)}), 500
 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860)
-
-
-# # Define the Gradio interface
-# interface = gr.Interface(
-#     fn=predict_text,
-#     inputs=[
-#         # gr.Image(type="pil", label="Image Input"),  # Image input with label
-#         gr.Textbox(label="Text Input")  # Textbox input with label
-#     ],
-#     outputs=gr.Textbox(label="Generated Response"),  # Output with a more descriptive label
-#     title="Llama 3.2 11B Vision Instruct Demo",  # Title of the interface
-#     description="This demo uses Meta's Llama 3.2 11B Vision model to generate responses based on an image and text input.",  # Short description
-#     theme="compact"  # Using a compact theme for a cleaner look
-# )
-
-# # Launch the interface
-# interface.launch(debug=True)
+    app.run(host="0.0.0.0", port=7860)