Update app.py

app.py

@@ -1,41 +1,50 @@
+from fastapi import FastAPI, HTTPException, UploadFile, File
+from transformers import pipeline
 import os
-from flask import Flask, request, jsonify
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import yaml
-
-# Load environment variables
-api_key = os.getenv('HF_API_KEY')
-model_path = os.getenv('MODEL_PATH')
-
-# Initialize Flask app
-app = Flask(__name__)
-
-# Load configuration
-with open('config.yaml', 'r') as file:
-    config = yaml.safe_load(file)
-
-# Load the model and tokenizer
-def load_model():
-    tokenizer = AutoTokenizer.from_pretrained(model_path)
-    model = AutoModelForCausalLM.from_pretrained(model_path)
-    return model, tokenizer
-
-model, tokenizer = load_model()
-
-def generate_text(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(inputs['input_ids'])
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-@app.route('/generate', methods=['POST'])
-def generate():
-    data = request.get_json()
-    prompt = data.get('prompt')
-    if prompt:
-        response_text = generate_text(prompt)
-        return jsonify({"response": response_text})
-    else:
-        return jsonify({"error": "No prompt provided"})
-
-if __name__ == '__main__':
-    app.run(debug=True)
+from typing import List
+from pydantic import BaseModel
+
+app = FastAPI()
+
+# Initialize the AI model
+model_name = "gpt-3.5-turbo"
+generator = pipeline('text-generation', model=model_name)
+
+class FileUpdate(BaseModel):
+    filename: str
+    content: str
+
+@app.post("/generate")
+def generate_text(prompt: str):
+    try:
+        result = generator(prompt, max_length=100)
+        return {"response": result[0]['generated_text']}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/upload")
+async def upload_file(file: UploadFile = File(...)):
+    try:
+        with open(f"./files/{file.filename}", "wb") as f:
+            content = await file.read()
+            f.write(content)
+        return {"filename": file.filename}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/update")
+def update_file(file_update: FileUpdate):
+    try:
+        with open(f"./files/{file_update.filename}", "w") as f:
+            f.write(file_update.content)
+        return {"status": "success"}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/files")
+def list_files():
+    try:
+        files = os.listdir("./files")
+        return {"files": files}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))