Upload 4 files

app.py

@@ -1,9 +1,7 @@
 import logging
 from fastapi import FastAPI, HTTPException
-from transformers import AutoModelForCausalLM, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from peft import PeftModel, PeftConfig
-from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
-from mistral_common.client import MistralChain
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -16,11 +14,10 @@ app = FastAPI()
 model = None
 tokenizer = None
 pipe = None
-mistral_chain = None
 
 @app.on_event("startup")
 async def load_model():
-    global model, tokenizer, pipe, mistral_chain
+    global model, tokenizer, pipe
     
     try:
         logger.info("Loading PEFT configuration...")
@@ -33,10 +30,7 @@ async def load_model():
         model = PeftModel.from_pretrained(base_model, "frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
 
         logger.info("Loading tokenizer...")
-        tokenizer = MistralTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
-
-        logger.info("Creating MistralChain...")
-        mistral_chain = MistralChain(model, tokenizer)
+        tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
 
         logger.info("Creating pipeline...")
         pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
@@ -55,12 +49,12 @@ def home():
 
 @app.get("/generate")
 async def generate(text: str):
-    if not mistral_chain:
+    if not pipe:
         raise HTTPException(status_code=503, detail="Model not loaded")
     
     try:
-        output = mistral_chain.generate(text, max_tokens=100)
-        return {"output": output}
+        output = pipe(text, max_length=100, num_return_sequences=1)
+        return {"output": output[0]['generated_text']}
     except Exception as e:
         logger.error(f"Error during text generation: {e}")
         raise HTTPException(status_code=500, detail=f"Error during text generation: {str(e)}")