Update app.py

app.py

@@ -10,6 +10,9 @@ from datetime import datetime
 import sqlite3
 from contextlib import contextmanager
 from werkzeug.utils import secure_filename
+from huggingface_hub import InferenceClient
+from langchain.llms.base import LLM
+from typing import Optional, List, Any
 
 # Initialize Flask application
 app = Flask(__name__)
@@ -64,6 +67,36 @@ prompt = PromptTemplate(
     template=prompt_template
 )
 
+class CustomHuggingFaceInference(LLM):
+    client: Any
+    model: str
+    
+    def __init__(self, token: str):
+        super().__init__()
+        self.client = InferenceClient(token=token)
+        self.model = "mistralai/Mistral-7B-Instruct-v0.1"
+    
+    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+        response = self.client.text_generation(
+            prompt,
+            model=self.model,
+            max_new_tokens=512,
+            temperature=0.7,
+            top_p=0.95,
+            repetition_penalty=1.15
+        )
+        return response
+
+    @property
+    def _identifying_params(self):
+        return {"model": self.model}
+    
+    @property
+    def _llm_type(self):
+        return "custom_huggingface"
+
+        
+
 @contextmanager
 def get_db_connection():
     """Context manager for database connections"""
@@ -117,19 +150,8 @@ def initialize_llm():
             raise ValueError("HUGGINGFACE_API_TOKEN not found in environment variables")
         print("API token found")
 
-        # Initialize with specific task parameters
-        llm = HuggingFaceHub(
-            repo_id="mistralai/Mistral-7B-Instruct-v0.1",
-            huggingfacehub_api_token=api_token,
-            task="text-generation",  # Specify the task
-            model_kwargs={
-                "temperature": 0.7,
-                "max_new_tokens": 512,
-                "top_p": 0.95,
-                "repetition_penalty": 1.15,
-                "return_full_text": False,
-            }
-        )
+        # Initialize with custom LLM
+        llm = CustomHuggingFaceInference(token=api_token)
         
         # Test the LLM
         print("Testing LLM with a simple prompt...")