Update app.py

app.py

@@ -18,6 +18,7 @@ import re
 from gradio_client import Client
 from simple_salesforce import Salesforce, SalesforceLogin
 from llama_index.llms.huggingface import HuggingFaceInferenceAPI
+from llama_index.llms.huggingface import HuggingFaceLLM
 
 
 # Define Pydantic model for incoming request body
@@ -28,11 +29,7 @@ class MessageRequest(BaseModel):
 #     model=repo_id,
 #     token=os.getenv("HF_TOKEN"),
 # )
-repo_id = "mistralai/Mistral-7B-Instruct-v0.1"
-llm_client = InferenceClient(
-    model=repo_id,
-    token=os.getenv("HF_TOKEN")
-)
+
 
 os.environ["HF_TOKEN"] = os.getenv("HF_TOKEN")
 username = os.getenv("username")
@@ -84,12 +81,14 @@ templates = Jinja2Templates(directory="static")
 #     max_new_tokens=512,
 #     generate_kwargs={"temperature": 0.1},
 # )
-Settings.llm = HuggingFaceInferenceAPI(
-    model_name="mistralai/Mistral-7B-Instruct-v0.1",
-    token=os.getenv("HF_TOKEN"),  # Your Hugging Face API token
-    context_window=4096,  # Mistral-7B’s context window
-    max_new_tokens=512,
-    generate_kwargs={"temperature": 0.1}
+# Configure Llama index settings
+Settings.llm = HuggingFaceLLM(
+    model_name="google/flan-t5-small",
+    tokenizer_name="google/flan-t5-small",
+    context_window=512,  # flan-t5-small has a max context window of 512 tokens
+    max_new_tokens=256,
+    generate_kwargs={"temperature": 0.1, "do_sample": True},
+    device_map="auto"  # Automatically use GPU if available, else CPU
 )
 Settings.embed_model = HuggingFaceEmbedding(
     model_name="BAAI/bge-small-en-v1.5"