Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import torch
 from fastapi import FastAPI
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Set environment variables
 os.environ["TRITON_DISABLE"] = "1"
@@ -18,8 +18,8 @@ os.environ["TORCH_HOME"] = "/tmp/hf_cache"
 # FastAPI app
 app = FastAPI()
 
-# Try loading a completely different model
-model_name = "facebook/opt-350m"  # Much smaller, more compatible model
+# Load your merged model
+model_name = "Suguru1846/counseling_model_merged"  # Your merged model
 tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir="/tmp/hf_cache")
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
@@ -31,7 +31,10 @@ model = AutoModelForCausalLM.from_pretrained(
 @app.post("/generate")
 async def generate_text(prompt: str, max_tokens: int = 50):
     try:
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        # Format prompt for Llama models
+        formatted_prompt = f"<s>[INST] {prompt} [/INST]"
+        
+        inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
         outputs = model.generate(
             **inputs, 
             max_new_tokens=max_tokens,
@@ -48,4 +51,4 @@ async def generate_text(prompt: str, max_tokens: int = 50):
 
 @app.get("/")
 async def root():
-    return {"message": "Model is Running"}
+    return {"message": "Your Custom Counseling Model is Running"}