Update app.py

app.py

@@ -1,25 +1,35 @@
 import os
+from fastapi import FastAPI
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+import torch
 
 # Set writable cache directory inside the container
 os.environ['SENTENCE_TRANSFORMERS_HOME'] = '/app/hf_home'
 os.environ['TRANSFORMERS_CACHE'] = '/app/hf_home'
 
-from fastapi import FastAPI
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
 # Ensure the directory exists
 os.makedirs(os.environ['TRANSFORMERS_CACHE'], exist_ok=True)
 
-# Load model
-model_name = "mynuddin/chatbot"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name).to("cpu")
+# Define base model and adapter model
+base_model_name = "facebook/opt-2.7b"
+adapter_name = "mynuddin/chatbot"
+
+# Load base model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+base_model = AutoModelForCausalLM.from_pretrained(base_model_name, torch_dtype=torch.float16)
+
+# Load PEFT adapter
+model = PeftModel.from_pretrained(base_model, adapter_name)
+model = model.to("cpu")  # Change to "cuda" if running on GPU
+model.eval()
 
 app = FastAPI()
 
 @app.post("/generate")
 def generate_text(prompt: str):
     inputs = tokenizer(prompt, return_tensors="pt")
-    output = model.generate(**inputs, max_length=128)
+    with torch.no_grad():
+        output = model.generate(**inputs, max_length=128)
     generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
-    return {"generated_query": generated_text}
+    return {"generated_query": generated_text}