Update app.py

app.py

@@ -1,81 +1,109 @@
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import json
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+import datetime
 
+# Initialize FastAPI
 app = FastAPI()
 
-# Model configuration
-MODEL_NAME = "nlptown/bert-base-multilingual-uncased-sentiment"
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Initialize sentiment analysis model
-sentiment_tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-sentiment_classifier = pipeline(
-    "sentiment-analysis",
-    model=MODEL_NAME,
-    tokenizer=sentiment_tokenizer,
-    device=DEVICE
+# Load model and tokenizer
+model_name = "Qwen/Qwen2.5-Coder-32B"
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map="auto",
+    trust_remote_code=True,
+    torch_dtype=torch.float16
 )
 
-# Initialize GPT-2 for text generation
-MODEL_NAME_LARGE = "gpt2-large"
-generation_tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME_LARGE)
-generation_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME_LARGE).to(DEVICE)
-
-class TextInput(BaseModel):
-    text: str
-
-class GenerationInput(BaseModel):
-    prompt: str
-    max_length: int = 100
-
-@app.post("/analyze-sentiment")
-async def analyze_sentiment(input_data: TextInput):
-    try:
-        result = sentiment_classifier(input_data.text)
-        return {
-            "sentiment": result[0]['label'],
-            "score": float(result[0]['score'])
+def format_chat_response(response_text, prompt_tokens, completion_tokens):
+    return {
+        "id": f"chatcmpl-{datetime.datetime.now().strftime('%Y%m%d%H%M%S')}",
+        "object": "chat.completion",
+        "created": int(datetime.datetime.now().timestamp()),
+        "model": model_name,
+        "choices": [{
+            "index": 0,
+            "message": {
+                "role": "assistant",
+                "content": response_text
+            },
+            "finish_reason": "stop"
+        }],
+        "usage": {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": prompt_tokens + completion_tokens
         }
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+    }
 
-@app.post("/generate-text")
-async def generate_text(input_data: GenerationInput):
+@app.post("/v1/chat/completions")
+async def chat_completion(request: Request):
     try:
-        inputs = generation_tokenizer(
-            input_data.prompt, 
-            return_tensors="pt"
-        ).to(DEVICE)
+        data = await request.json()
+        messages = data.get("messages", [])
+        
+        # Format messages for Qwen
+        conversation = []
+        for msg in messages:
+            conversation.append({
+                "role": msg["role"],
+                "content": msg["content"]
+            })
         
-        outputs = generation_model.generate(
-            inputs["input_ids"],
-            max_length=input_data.max_length,
-            num_return_sequences=1,
-            no_repeat_ngram_size=2,
-            pad_token_id=generation_tokenizer.eos_token_id
+        # Convert messages to model input format
+        prompt = tokenizer.apply_chat_template(
+            conversation,
+            tokenize=False,
+            add_generation_prompt=True
         )
         
-        generated_text = generation_tokenizer.decode(
-            outputs[0], 
-            skip_special_tokens=True
+        # Count prompt tokens
+        prompt_tokens = len(tokenizer.encode(prompt))
+        
+        # Generate response
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=data.get("max_tokens", 2048),
+            temperature=data.get("temperature", 0.7),
+            top_p=data.get("top_p", 0.95),
+            do_sample=True
         )
         
-        return {"generated_text": generated_text}
+        response_text = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+        completion_tokens = len(tokenizer.encode(response_text))
+        
+        return JSONResponse(
+            content=format_chat_response(response_text, prompt_tokens, completion_tokens)
+        )
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        return JSONResponse(
+            status_code=500,
+            content={"error": str(e)}
+        )
 
-@app.get("/health")
-async def health_check():
-    return {
-        "status": "healthy",
-        "sentiment_model": MODEL_NAME,
-        "generation_model": MODEL_NAME_LARGE,
-        "device": str(DEVICE)
-    }
+# Gradio interface for testing
+def chat_interface(message, history):
+    history = history or []
+    messages = [{"role": "user", "content": message}]
+    
+    # Add history to messages
+    for h in history:
+        messages.insert(0, {"role": "assistant" if i % 2 else "user", "content": h[1 if i % 2 else 0]} 
+                       for i in range(len(h)))
+    
+    response = chat_completion(Request({"messages": messages}))
+    return response.choices[0].message.content
+
+interface = gr.ChatInterface(
+    chat_interface,
+    title="Qwen2.5-Coder-32B Chat",
+    description="Chat with Qwen2.5-Coder-32B model. This Space also provides a /v1/chat/completions endpoint."
+)
 
-# Dodaj to na końcu pliku
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+# Mount both FastAPI and Gradio
+app = gr.mount_gradio_app(app, interface, path="/")