Update app.py

app.py

@@ -7,7 +7,7 @@ import torch
 from torchvision import transforms
 from torchvision.models.detection import fasterrcnn_resnet50_fpn
 from PIL import Image
-from transformers import pipeline
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
 from fastapi.responses import RedirectResponse
 import numpy as np
@@ -16,25 +16,19 @@ import easyocr
 # Initialize FastAPI
 app = FastAPI()
 
-# Load AI Model for Question Answering (DeepSeek-V2-Chat)
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-# Preload Hugging Face model
-model_name = "microsoft/phi-2"
+# Load AI Model for Question Answering (Mistral-7B)
+model_name = "mistralai/Mistral-7B"
 print(f"🔄 Loading model: {model_name}...")
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
 
-qa_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
+qa_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)
 
 # Load Pretrained Object Detection Model (Torchvision)
 from torchvision.models.detection import FasterRCNN_ResNet50_FPN_Weights
 weights = FasterRCNN_ResNet50_FPN_Weights.DEFAULT
-model = fasterrcnn_resnet50_fpn(weights=weights)
-model.eval()
-# Load Pretrained Object Detection Model (if needed)
-model = fasterrcnn_resnet50_fpn(pretrained=True)
-model.eval()
+object_detection_model = fasterrcnn_resnet50_fpn(weights=weights)
+object_detection_model.eval()
 
 # Initialize OCR Model (Lazy Load)
 reader = easyocr.Reader(["en"], gpu=True)
@@ -172,4 +166,4 @@ app = gr.mount_gradio_app(app, demo, path="/")
 
 @app.get("/")
 def home():
-    return RedirectResponse(url="/")
+    return RedirectResponse(url="/")