updated

app.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel
 import torch
 
@@ -31,27 +31,10 @@ class GenerateRequest(BaseModel):
 # **Use model.generate() instead of pipeline()**
 def generate_text_from_model(prompt: str):
     try:
-        input_data = tokenizer(
-            f"<s>[INST] {prompt} [/INST]", 
-            return_tensors="pt", 
-            padding=True, 
-            truncation=True
-        )
-        input_ids = input_data.input_ids.to(device)
-        attention_mask = input_data.attention_mask.to(device)
-        
-        output_ids = model.generate(
-            input_ids, 
-            max_length=512, 
-            pad_token_id=tokenizer.eos_token_id,
-            attention_mask=attention_mask
-        )
-        
+        input_ids = tokenizer(f"<s>[INST] {prompt} [/INST]", return_tensors="pt").input_ids.to(device)
+        output_ids = model.generate(input_ids, max_length=512)
         generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-        
-        # Extract only the assistant's response
-        response_text = generated_text.split("<|assistant|>\n")[-1].strip()
-        return response_text
+        return generated_text
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
@@ -64,5 +47,4 @@ async def root():
 @app.post("/generate/")
 async def generate_text(request: GenerateRequest):
     response = generate_text_from_model(request.prompt)
-    return {"response": response}
-   
+    return {"response": response}