Update main.py

main.py

@@ -1,64 +1,20 @@
-from flask import Flask, request, jsonify
+from fastapi import FastAPI
 from transformers import pipeline
-import torch
-import os
 
-app = Flask(__name__)
 
-# Check for MPS (Apple Silicon) support
-if torch.backends.mps.is_available():
-    device = torch.device("mps")
-elif torch.cuda.is_available():
-    device = torch.device("cuda")
-else:
-    device = torch.device("cpu")
-print(f"Using device: {device}")
+# NOTE - we configure docs_url to serve the interactive Docs at the root path
+# of the app. This way, we can use the docs as a landing page for the app on Spaces.
+app = FastAPI(docs_url="/")
 
-# Load the Hugging Face models
-model_name = "./Qwen2.5-1.5B-Instruct-Local"  # or any other model suitable for your needs
-chatbot_model = pipeline("text-generation", model=model_name, device=device, batch_size=8)
-# sentiment_model = pipeline("sentiment-analysis", device=device)
-
-# Store conversation history
-conversation_history = []
-
-def generate_response(user_input, sentiment):
-    # if sentiment['label'] == 'NEGATIVE':
-    #     prompt = f"The user is upset, respond with empathy and support: {user_input}"
-    # else:
-    #     prompt = f"Respond to the following query: {user_input}"
-
-    # Append the user prompt to the conversation history
-    conversation_history.append({"role": "user", "content": prompt})
-
-    # Generate response from the model
-    result = chatbot_model(conversation_history, num_return_sequences=1, max_new_tokens=250)
-
-    # Append the assistant's response to the conversation history
-    conversation_history.append({"role": "assistant", "content": result[0]['generated_text'][-1]['content'] })
-
-    return result[0]['generated_text'][-1]['content'] 
-
-
-@app.route('/chat', methods=['POST'])
-def chat():
-    data = request.json
-    user_input = data.get('message', '')
-
-    if not user_input:
-        return jsonify({'error': 'Please provide a message.'}), 400
-
-    # sentiment = sentiment_model(user_input)[0]
-    response = generate_response(user_input, 'POSITIVE')
-    
-    with open(file_path, 'w') as file:
-        file.write("\nConversation History:\n")
-        for item in conversation_history:
-            file.write(f"{item}\n")
-    return jsonify({'response': response , 'sentiment': sentiment})
-
-@app.route('/history', methods=['GET'])
-def get_history():
-    return jsonify(conversation_history)
+pipe = pipeline("text2text-generation", model="google/flan-t5-small")
 
 
+@app.get("/generate")
+def generate(text: str):
+    """
+    Using the text2text-generation pipeline from `transformers`, generate text
+    from the given input text. The model used is `google/flan-t5-small`, which
+    can be found [here](https://huggingface.co/google/flan-t5-small).
+    """
+    output = pipe(text)
+    return {"output": output[0]["generated_text"]}