Upload main.py

main.py

@@ -5,7 +5,7 @@ import os
 import fitz  # PyMuPDF for PDF handling
 import torch
 import json
-from transformers import pipeline, GPT2Tokenizer, GPT2LMHeadModel
+from transformers import pipeline, GPT2Tokenizer, AutoModelForCausalLM
 import logging
 
 # Initialize FastAPI app
@@ -15,11 +15,15 @@ app = FastAPI()
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Initialize the GPT2-medium pipeline
+# Load GPT-Neo-2.7B tokenizer and model
+tokenizer = GPT2Tokenizer.from_pretrained("EleutherAI/gpt-neo-2.7B")
+model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-2.7B")
+
+# Initialize the text-generation pipeline
 qa_pipeline = pipeline(
     "text-generation",
-    model="gpt2-medium",  # Switching to GPT2-medium
-    tokenizer="gpt2-medium",
+    model=model,
+    tokenizer=tokenizer,
     device=-1  # Use CPU
 )
 
@@ -39,62 +43,81 @@ class ExtractedInfo(BaseModel):
 
 def extract_text_from_pdf(pdf_path: str) -> str:
     """Extracts text from a PDF file."""
-    with fitz.open(pdf_path) as doc:
-        text = "".join([page.get_text() for page in doc])
-    if not text.strip():
-        raise HTTPException(status_code=400, detail="PDF contains no extractable text.")
-    return text
-
-def chunk_text(text: str, max_tokens: int = 900) -> list:
-    """Splits the text into chunks that fit the model token limit."""
-    tokens = qa_pipeline.tokenizer.encode(text)
-    chunks = [
-        qa_pipeline.tokenizer.decode(tokens[i:i + max_tokens])
+    try:
+        with fitz.open(pdf_path) as doc:
+            text = "".join([page.get_text() for page in doc])
+        if not text.strip():
+            raise ValueError("PDF contains no extractable text.")
+        return text
+    except Exception as e:
+        logger.error(f"Error extracting text from PDF: {e}")
+        raise HTTPException(status_code=400, detail="Failed to extract text from the PDF.")
+
+def chunk_text(text: str, max_tokens: int = 1800) -> list:
+    """Chunks text using the tokenizer to fit within GPT-Neo's 2048-token limit."""
+    tokens = tokenizer.encode(text)
+    return [
+        tokenizer.decode(tokens[i : i + max_tokens])
         for i in range(0, len(tokens), max_tokens)
     ]
-    return chunks
+
+def generate_chunk_output(chunk: str) -> dict:
+    """Generates structured output for a single chunk."""
+    prompt = f"""
+    Extract the following information from the resume in JSON format:
+    {{
+        "work_experience": "<Summarized single work experience>",
+        "education": {{
+            "degree": "<Degree obtained>",
+            "university": "<University attended>",
+            "graduation_year": "<Year of graduation>"
+        }},
+        "professional_course_detail": "<Details of professional courses completed>",
+        "software_usage": "<List of software tools used>",
+        "safety_course_detail": "<Safety courses completed>",
+        "hse_description": "<HSE (Health, Safety, Environment) practices>",
+        "good_conduct_certificate": "<Details of good conduct certificate>"
+    }}
+    Resume text:
+    {chunk}
+    """
+
+    response = qa_pipeline(prompt, max_new_tokens=500, temperature=0.7)
+    generated_text = response[0]["generated_text"]
+
+    # Extract JSON from the generated text
+    json_start = generated_text.find("{")
+    json_end = generated_text.rfind("}") + 1
+    if json_start != -1 and json_end != -1:
+        json_str = generated_text[json_start:json_end]
+        return json.loads(json_str)
+    else:
+        logger.warning("Invalid JSON output from model.")
+        return {}  # Return empty dict if parsing fails
 
 def generate_structured_output(text: str) -> dict:
-    """Generates structured output from the text using GPT2-medium."""
+    """Generates structured output from multiple chunks."""
     chunks = chunk_text(text)
-
-    # Collect results from each chunk
-    generated_text = ""
+    merged_output = {
+        "work_experience": "",
+        "education": {"degree": "", "university": "", "graduation_year": ""},
+        "professional_course_detail": "",
+        "software_usage": "",
+        "safety_course_detail": "",
+        "hse_description": "",
+        "good_conduct_certificate": ""
+    }
+
+    # Process each chunk and merge results
     for chunk in chunks:
-        prompt = f"""
-        Extract the following information from the resume in JSON format:
-        {{
-            "work_experience": "<Summarized single work experience>",
-            "education": {{
-                "degree": "<Degree obtained>",
-                "university": "<University attended>",
-                "graduation_year": "<Year of graduation>"
-            }},
-            "professional_course_detail": "<Details of professional courses completed>",
-            "software_usage": "<List of software tools used>",
-            "safety_course_detail": "<Safety courses completed>",
-            "hse_description": "<HSE (Health, Safety, Environment) practices>",
-            "good_conduct_certificate": "<Details of good conduct certificate>"
-        }}
-        Resume text:
-        {chunk}
-        """
-
-        response = qa_pipeline(prompt, max_new_tokens=300, temperature=0.7)
-        generated_text += response[0]["generated_text"]
-
-    # Extract JSON from the generated text
-    try:
-        json_start = generated_text.find("{")
-        json_end = generated_text.rfind("}") + 1
-        if json_start != -1 and json_end != -1:
-            json_str = generated_text[json_start:json_end]
-            return json.loads(json_str)
-        else:
-            raise ValueError("No valid JSON found in the model output")
-    except Exception as e:
-        logger.error(f"Error generating structured output: {e}")
-        raise HTTPException(status_code=500, detail="Failed to generate structured output.")
+        chunk_output = generate_chunk_output(chunk)
+        for key, value in chunk_output.items():
+            if isinstance(value, dict):
+                merged_output[key].update(value)
+            elif not merged_output[key]:  # Only fill if not already populated
+                merged_output[key] = value
+
+    return merged_output
 
 @app.post("/process_cv/", response_model=ExtractedInfo)
 async def process_cv(background_tasks: BackgroundTasks, file: UploadFile = File(...)):